geo/server/internal/tenant/app/runtime_config.go

package app

import (
	"strings"
	"time"

	"github.com/geo-platform/tenant-api/internal/shared/config"
)

const defaultArticleGenerationTimeout = 8 * time.Minute

type runtimeConfigProvider = config.Provider

type generationRuntimeConfig struct {
	StreamEnabled   bool
	ArticleTimeout  time.Duration
	MaxOutputTokens int64
}

type knowledgeRuntimeConfig struct {
	ChunkSize          int
	ChunkOverlap       int
	EmbeddingBatchSize int
	RecallLimit        int
	RerankTopN         int
	ArkBaseURL         string
	ArkAPIKey          string
	URLMarkdownModel   string
	BrowserFetch       config.BrowserFetchConfig
}

type generationStreamDisabledConfigProvider struct {
	base config.Provider
}

func NewGenerationStreamDisabledConfigProvider(base config.Provider) config.Provider {
	return generationStreamDisabledConfigProvider{base: base}
}

func (p generationStreamDisabledConfigProvider) Current() *config.Config {
	if p.base == nil {
		return nil
	}
	cfg := p.base.Current()
	if cfg == nil {
		return nil
	}
	next := *cfg
	next.Generation.StreamEnabled = false
	return &next
}

func generationRuntimeFromConfig(generation config.GenerationConfig, llmCfg config.LLMConfig) generationRuntimeConfig {
	articleTimeout := generation.ArticleTimeout
	if articleTimeout <= 0 {
		articleTimeout = defaultArticleGenerationTimeout
	}
	return generationRuntimeConfig{
		StreamEnabled:   generation.StreamEnabled,
		ArticleTimeout:  articleTimeout,
		MaxOutputTokens: llmCfg.MaxOutputTokens,
	}
}

func knowledgeRuntimeFromConfig(retrievalCfg config.RetrievalConfig, llmCfg config.LLMConfig) knowledgeRuntimeConfig {
	chunkSize := retrievalCfg.ChunkSize
	if chunkSize <= 0 {
		chunkSize = 900
	}
	chunkOverlap := retrievalCfg.ChunkOverlap
	if chunkOverlap < 0 {
		chunkOverlap = 120
	}
	batchSize := retrievalCfg.EmbeddingBatchSize
	if batchSize <= 0 {
		batchSize = 16
	}
	recallLimit := retrievalCfg.RecallLimit
	if recallLimit <= 0 {
		recallLimit = 12
	}
	rerankTopN := retrievalCfg.RerankTopN
	if rerankTopN <= 0 {
		rerankTopN = 6
	}

	model := strings.TrimSpace(llmCfg.KnowledgeURLModel)
	if model == "" {
		model = strings.TrimSpace(llmCfg.Model)
	}
	baseURL := strings.TrimSpace(llmCfg.BaseURL)
	if baseURL == "" {
		baseURL = defaultKnowledgeArkBaseURL
	}

	return knowledgeRuntimeConfig{
		ChunkSize:          chunkSize,
		ChunkOverlap:       chunkOverlap,
		EmbeddingBatchSize: batchSize,
		RecallLimit:        recallLimit,
		RerankTopN:         rerankTopN,
		ArkBaseURL:         baseURL,
		ArkAPIKey:          strings.TrimSpace(llmCfg.APIKey),
		URLMarkdownModel:   model,
	}
}

func knowledgeRuntimeFromFullConfig(cfg *config.Config) knowledgeRuntimeConfig {
	if cfg == nil {
		return knowledgeRuntimeFromConfig(config.RetrievalConfig{}, config.LLMConfig{})
	}
	runtime := knowledgeRuntimeFromConfig(cfg.Retrieval, cfg.LLM)
	browserFetch := cfg.BrowserFetch
	config.NormalizeBrowserFetchConfig(&browserFetch)
	runtime.BrowserFetch = browserFetch
	return runtime
}