feat: add document and embedding pipeline with hybrid search

Introduce a document system with chunking, embedding generation (Ollama/OpenAI),
and async pipeline for semantic search. Refactor hybrid search from entity package
into dedicated embedding package, replace concrete ChunkRepository with
EmbeddingRepository, and add document REST API and MCP tools.

Author: ravisuhag

cli/embed.go

@@ -0,0 +1,233 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"strings"
+
+	"github.com/MakeNowJust/heredoc"
+	"github.com/raystack/compass/core/chunking"
+	"github.com/raystack/compass/core/document"
+	"github.com/raystack/compass/core/embedding"
+	"github.com/raystack/compass/core/entity"
+	"github.com/raystack/compass/core/namespace"
+	"github.com/raystack/compass/internal/config"
+	compassserver "github.com/raystack/compass/internal/server"
+	"github.com/raystack/compass/store/postgres"
+	"github.com/spf13/cobra"
+)
+
+func embedCommand(cfg *config.Config) *cobra.Command {
+	var embedType string
+	var batchSize int
+
+	cmd := &cobra.Command{
+		Use:   "embed",
+		Short: "Generate embeddings for entities and documents",
+		Long:  "Batch generate embeddings for all entities and/or documents. Requires an embedding provider configured.",
+		Example: heredoc.Doc(`
+			$ compass embed
+			$ compass embed --type entity
+			$ compass embed --type document
+			$ compass embed --batch-size 50
+		`),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			compassserver.InitLogger(cfg.LogLevel)
+			return runEmbed(cmd.Context(), cfg, embedType, batchSize)
+		},
+	}
+
+	cmd.Flags().StringVar(&embedType, "type", "all", "Type to embed: entity, document, or all")
+	cmd.Flags().IntVar(&batchSize, "batch-size", 100, "Number of items to process per batch")
+
+	return cmd
+}
+
+func runEmbed(ctx context.Context, cfg *config.Config, embedType string, batchSize int) error {
+	if !cfg.Embedding.Enabled {
+		return fmt.Errorf("embedding is not enabled in config")
+	}
+
+	// Init embedding provider
+	provider, err := initProvider(cfg.Embedding)
+	if err != nil {
+		return err
+	}
+	slog.Info("embedding provider initialized", "provider", provider.Name())
+
+	// Init postgres
+	pgClient, err := postgres.NewClient(cfg.DB)
+	if err != nil {
+		return fmt.Errorf("connect to postgres: %w", err)
+	}
+	defer pgClient.Close()
+
+	embeddingRepo, err := postgres.NewEmbeddingRepository(pgClient)
+	if err != nil {
+		return err
+	}
+
+	ns := namespace.DefaultNamespace
+
+	if embedType == "all" || embedType == "entity" {
+		entityRepo, err := postgres.NewEntityRepository(pgClient)
+		if err != nil {
+			return err
+		}
+		if err := embedEntities(ctx, entityRepo, embeddingRepo, provider, ns, batchSize, cfg.Embedding); err != nil {
+			return fmt.Errorf("embed entities: %w", err)
+		}
+	}
+
+	if embedType == "all" || embedType == "document" {
+		docRepo, err := postgres.NewDocumentRepository(pgClient)
+		if err != nil {
+			return err
+		}
+		if err := embedDocuments(ctx, docRepo, embeddingRepo, provider, ns, batchSize, cfg.Embedding); err != nil {
+			return fmt.Errorf("embed documents: %w", err)
+		}
+	}
+
+	slog.Info("embedding complete")
+	return nil
+}
+
+func embedEntities(ctx context.Context, entityRepo entity.Repository, embeddingRepo embedding.Repository,
+	provider embedding.Provider, ns *namespace.Namespace, batchSize int, cfg config.EmbeddingConfig) error {
+
+	offset := 0
+	total := 0
+
+	for {
+		entities, err := entityRepo.GetAll(ctx, ns, entity.Filter{Size: batchSize, Offset: offset})
+		if err != nil {
+			return err
+		}
+		if len(entities) == 0 {
+			break
+		}
+
+		for _, ent := range entities {
+			chunks := chunking.SerializeEntity(ent)
+			if len(chunks) == 0 {
+				continue
+			}
+
+			text := chunks[0].Context + "\n\n" + chunks[0].Content
+			vec, err := provider.Embed(ctx, text)
+			if err != nil {
+				slog.Error("failed to embed entity", "urn", ent.URN, "error", err)
+				continue
+			}
+
+			embs := []embedding.Embedding{{
+				EntityURN:   ent.URN,
+				ContentID:   ent.ID,
+				ContentType: "entity",
+				Content:     chunks[0].Content,
+				Context:     chunks[0].Context,
+				Vector:      vec,
+				Position:    0,
+				Heading:     chunks[0].Heading,
+				TokenCount:  chunking.EstimateTokens(chunks[0].Content),
+			}}
+
+			if err := embeddingRepo.UpsertBatch(ctx, ns, embs); err != nil {
+				slog.Error("failed to store entity embedding", "urn", ent.URN, "error", err)
+				continue
+			}
+			total++
+		}
+
+		slog.Info("embedded entities", "batch", offset/batchSize+1, "count", len(entities), "total", total)
+		offset += batchSize
+
+		if len(entities) < batchSize {
+			break
+		}
+	}
+
+	slog.Info("entity embedding complete", "total", total)
+	return nil
+}
+
+func embedDocuments(ctx context.Context, docRepo document.Repository, embeddingRepo embedding.Repository,
+	provider embedding.Provider, ns *namespace.Namespace, batchSize int, cfg config.EmbeddingConfig) error {
+
+	docs, err := docRepo.GetAll(ctx, ns, document.Filter{Size: 10000})
+	if err != nil {
+		return err
+	}
+
+	total := 0
+	for _, doc := range docs {
+		chunks := chunking.SplitDocument(doc.Title, doc.Body, chunking.Options{
+			MaxTokens: cfg.MaxTokens,
+			Overlap:   cfg.Overlap,
+			Title:     doc.Title,
+		})
+		if len(chunks) == 0 {
+			continue
+		}
+
+		// Prepare texts for batch embedding
+		texts := make([]string, len(chunks))
+		for i, c := range chunks {
+			if c.Context != "" {
+				texts[i] = c.Context + "\n\n" + c.Content
+			} else {
+				texts[i] = c.Content
+			}
+		}
+
+		vectors, err := provider.EmbedBatch(ctx, texts)
+		if err != nil {
+			slog.Error("failed to embed document", "title", doc.Title, "entity_urn", doc.EntityURN, "error", err)
+			continue
+		}
+
+		embs := make([]embedding.Embedding, len(chunks))
+		for i, c := range chunks {
+			var vec []float32
+			if i < len(vectors) {
+				vec = vectors[i]
+			}
+			embs[i] = embedding.Embedding{
+				EntityURN:   doc.EntityURN,
+				ContentID:   doc.ID,
+				ContentType: "document",
+				Content:     c.Content,
+				Context:     c.Context,
+				Vector:      vec,
+				Position:    c.Position,
+				Heading:     c.Heading,
+				TokenCount:  chunking.EstimateTokens(c.Content),
+			}
+		}
+
+		if err := embeddingRepo.UpsertBatch(ctx, ns, embs); err != nil {
+			slog.Error("failed to store document embeddings", "title", doc.Title, "error", err)
+			continue
+		}
+		total += len(embs)
+	}
+
+	slog.Info("document embedding complete", "total_chunks", total, "documents", len(docs))
+	return nil
+}
+
+func initProvider(cfg config.EmbeddingConfig) (embedding.Provider, error) {
+	switch strings.ToLower(cfg.Provider) {
+	case "openai":
+		if cfg.OpenAI.APIKey == "" {
+			return nil, fmt.Errorf("openai api_key is required")
+		}
+		return embedding.NewOpenAI(cfg.OpenAI), nil
+	case "ollama", "":
+		return embedding.NewOllama(cfg.Ollama), nil
+	default:
+		return nil, fmt.Errorf("unknown embedding provider: %s", cfg.Provider)
+	}
+}

cli/root.go

@@ -57,6 +57,7 @@ func New(cliConfig *config.Config) *cobra.Command {
 		configCommand(cliConfig),
 		namespacesCommand(cliConfig),
 		entitiesCommand(cliConfig),
+		embedCommand(cliConfig),
 		versionCmd(),
 	)
 

core/chunking/chunker.go

@@ -0,0 +1,46 @@
+package chunking
+
+// Chunk is a text fragment ready for embedding.
+type Chunk struct {
+	Content  string // the text to embed
+	Context  string // contextual prefix for embedding quality
+	Heading  string // section heading
+	Position int    // ordering within parent
+}
+
+// Options configures chunking behavior.
+type Options struct {
+	MaxTokens int    // target chunk size in tokens (default 512)
+	Overlap   int    // overlap tokens between adjacent chunks (default 50)
+	Title     string // parent title for contextual prefix
+}
+
+func (o Options) maxTokens() int {
+	if o.MaxTokens <= 0 {
+		return 512
+	}
+	return o.MaxTokens
+}
+
+func (o Options) overlap() int {
+	if o.Overlap <= 0 {
+		return 50
+	}
+	return o.Overlap
+}
+
+// EstimateTokens estimates the token count for a text string.
+// Uses a simple heuristic: ~1.3 tokens per word on average.
+func EstimateTokens(text string) int {
+	words := 0
+	inWord := false
+	for _, r := range text {
+		if r == ' ' || r == '\n' || r == '\t' || r == '\r' {
+			inWord = false
+		} else if !inWord {
+			inWord = true
+			words++
+		}
+	}
+	return (words * 4) / 3
+}

core/chunking/serializer.go

@@ -0,0 +1,101 @@
+package chunking
+
+import (
+	"fmt"
+	"sort"
+	"strings"
+
+	"github.com/raystack/compass/core/entity"
+)
+
+// SerializeEntity converts an entity's metadata into readable text for embedding.
+// Most entities produce a single chunk.
+func SerializeEntity(ent entity.Entity) []Chunk {
+	var b strings.Builder
+
+	fmt.Fprintf(&b, "Name: %s\n", ent.Name)
+	fmt.Fprintf(&b, "Type: %s\n", ent.Type)
+	fmt.Fprintf(&b, "URN: %s\n", ent.URN)
+	if ent.Source != "" {
+		fmt.Fprintf(&b, "Source: %s\n", ent.Source)
+	}
+	if ent.Description != "" {
+		fmt.Fprintf(&b, "Description: %s\n", ent.Description)
+	}
+
+	if len(ent.Properties) > 0 {
+		b.WriteString("Properties:\n")
+		serializeProperties(&b, ent.Properties, "  ")
+	}
+
+	context := fmt.Sprintf("%s: %s (%s)", ent.Type, ent.Name, ent.URN)
+
+	return []Chunk{{
+		Content:  b.String(),
+		Context:  context,
+		Heading:  string(ent.Type) + ": " + ent.Name,
+		Position: 0,
+	}}
+}
+
+func serializeProperties(b *strings.Builder, props map[string]interface{}, indent string) {
+	// Sort keys for deterministic output
+	keys := make([]string, 0, len(props))
+	for k := range props {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+
+	for _, k := range keys {
+		v := props[k]
+		switch val := v.(type) {
+		case map[string]interface{}:
+			fmt.Fprintf(b, "%s%s:\n", indent, k)
+			serializeProperties(b, val, indent+"  ")
+		case []interface{}:
+			serializeSlice(b, k, val, indent)
+		default:
+			fmt.Fprintf(b, "%s%s: %v\n", indent, k, v)
+		}
+	}
+}
+
+func serializeSlice(b *strings.Builder, key string, items []interface{}, indent string) {
+	if len(items) == 0 {
+		return
+	}
+
+	// Check if items are simple strings/numbers
+	allSimple := true
+	for _, item := range items {
+		switch item.(type) {
+		case string, float64, int, bool:
+		default:
+			allSimple = false
+		}
+	}
+
+	if allSimple {
+		vals := make([]string, len(items))
+		for i, item := range items {
+			vals[i] = fmt.Sprintf("%v", item)
+		}
+		fmt.Fprintf(b, "%s%s: %s\n", indent, key, strings.Join(vals, ", "))
+		return
+	}
+
+	// Complex items (e.g., columns)
+	fmt.Fprintf(b, "%s%s:\n", indent, key)
+	for _, item := range items {
+		if m, ok := item.(map[string]interface{}); ok {
+			// Try to format as "name (type)" for column-like structures
+			name, _ := m["name"].(string)
+			typ, _ := m["type"].(string)
+			if name != "" && typ != "" {
+				fmt.Fprintf(b, "%s  - %s (%s)\n", indent, name, typ)
+				continue
+			}
+		}
+		fmt.Fprintf(b, "%s  - %v\n", indent, item)
+	}
+}