yao/kb/providers/factory/factory.go

package factory

import (
	"fmt"

	"github.com/yaoapp/gou/graphrag/types"
	kbtypes "github.com/yaoapp/yao/kb/types"
)

// ProviderType is a type for provider types
type ProviderType string

const (
	// ProviderTypeChunking is a type for chunking providers
	ProviderTypeChunking ProviderType = "chunking"
	// ProviderTypeConverter is a type for converter providers
	ProviderTypeConverter ProviderType = "converter"
	// ProviderTypeEmbedding is a type for embedding providers
	ProviderTypeEmbedding ProviderType = "embedding"
	// ProviderTypeExtractor is a type for extractor providers
	ProviderTypeExtractor ProviderType = "extractor"
	// ProviderTypeFetcher is a type for fetcher providers
	ProviderTypeFetcher ProviderType = "fetcher"
)

// DetectMatch is a match for auto detect
type DetectMatch struct {
	ID       string
	Priority int
}

// Chunkings is a map of chunking providers
var Chunkings = map[string]Chunking{}

// Converters is a map of converter providers
var Converters = map[string]Converter{}

// Embeddings is a map of embedding providers
var Embeddings = map[string]Embedding{}

// Extractors is a map of extractor providers
var Extractors = map[string]Extractor{}

// Fetchers is a map of fetcher providers
var Fetchers = map[string]Fetcher{}

// === Chunking API ===

// MakeChunking creates a new chunking provider
func MakeChunking(id string, option *kbtypes.ProviderOption) (types.Chunking, error) {
	chunking, ok := Chunkings[id]
	if !ok {
		return nil, fmt.Errorf("chunking provider %s not found", id)
	}
	return chunking.Make(option)
}

// ChunkingOptions returns the options for a chunking provider
func ChunkingOptions(id string, option *kbtypes.ProviderOption) (*types.ChunkingOptions, error) {
	chunking, ok := Chunkings[id]
	if !ok {
		return nil, fmt.Errorf("chunking provider %s not found", id)
	}
	return chunking.Options(option)
}

// === Converter API ===

// MakeConverter creates a new converter provider
func MakeConverter(id string, option *kbtypes.ProviderOption) (types.Converter, error) {
	converter, ok := Converters[id]
	if !ok {
		return nil, fmt.Errorf("converter provider %s not found", id)
	}
	return converter.Make(option)
}

// AutoDetectConverter detects the converter based on the filename and content types
// return matched, id, error
func AutoDetectConverter(filename, contentType string) (bool, string, error) {
	var highestPriority int = 0
	var highestID string = ""
	for id, converter := range Converters {
		ok, priority, err := converter.AutoDetect(filename, contentType)
		if err != nil {
			continue
		}
		if ok && priority > highestPriority {
			highestPriority = priority
			highestID = id
		}
	}
	return highestID != "", highestID, nil
}

// === Embedding API ===

// MakeEmbedding creates a new embedding provider
func MakeEmbedding(id string, option *kbtypes.ProviderOption) (types.Embedding, error) {
	embedding, ok := Embeddings[id]
	if !ok {
		return nil, fmt.Errorf("embedding provider %s not found", id)
	}
	return embedding.Make(option)
}

// === Extractor API ===

// MakeExtractor creates a new extractor provider
func MakeExtractor(id string, option *kbtypes.ProviderOption) (types.Extraction, error) {
	extractor, ok := Extractors[id]
	if !ok {
		return nil, fmt.Errorf("extractor provider %s not found", id)
	}
	return extractor.Make(option)
}

// === Fetcher API ===

// MakeFetcher creates a new fetcher provider
func MakeFetcher(id string, option *kbtypes.ProviderOption) (types.Fetcher, error) {
	fetcher, ok := Fetchers[id]
	if !ok {
		return nil, fmt.Errorf("fetcher provider %s not found", id)
	}
	return fetcher.Make(option)
}

// === Schema API ===

// GetSchema returns the schema for a provider
func GetSchema(typ ProviderType, provider *kbtypes.Provider) (*kbtypes.ProviderSchema, error) {
	var schema Schema = nil
	var exists bool = false
	switch typ {
	case ProviderTypeChunking:
		schema, exists = Chunkings[provider.ID]
	case ProviderTypeConverter:
		schema, exists = Converters[provider.ID]
	case ProviderTypeEmbedding:
		schema, exists = Embeddings[provider.ID]
	case ProviderTypeExtractor:
		schema, exists = Extractors[provider.ID]
	case ProviderTypeFetcher:
		schema, exists = Fetchers[provider.ID]
	}
	if !exists {
		return nil, fmt.Errorf("%s provider %s not found", typ, provider.ID)
	}
	return schema.Schema(provider)
}