geo/server/internal/tenant/prompts/loader.go

package prompts

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"gopkg.in/yaml.v3"
)

const promptsConfigEnv = "PROMPTS_CONFIG_FILE"

type promptsConfig struct {
	Runtime           runtimePromptsConfig                   `yaml:"runtime"`
	PlatformTemplates map[string]platformTemplatePromptEntry `yaml:"platform_templates"`
}

type runtimePromptsConfig struct {
	DefaultGenerationBasePromptTemplate    string                         `yaml:"default_generation_base_prompt_template"`
	PromptContextHeading                   string                         `yaml:"prompt_context_heading"`
	GenerationWritingRequirementsSection   string                         `yaml:"generation_writing_requirements_section"`
	GenerationTemplateSpecificRulesHeading string                         `yaml:"generation_template_specific_rules_heading"`
	GenerationLengthGuidanceHeading        string                         `yaml:"generation_length_guidance_heading"`
	TopXBrandPriorityRulesTemplate         string                         `yaml:"top_x_brand_priority_rules_template"`
	EnglishLengthGuidanceTemplate          string                         `yaml:"english_length_guidance_template"`
	ChineseLengthGuidanceTemplate          string                         `yaml:"chinese_length_guidance_template"`
	ContextLabels                          map[string]string              `yaml:"context_labels"`
	TemplateContextWithJSONExampleTemplate string                         `yaml:"template_context_with_json_example_template"`
	JSONOutputExampleHeading               string                         `yaml:"json_output_example_heading"`
	AnalyzeOutputExample                   string                         `yaml:"analyze_output_example"`
	AnalyzeFallbackPromptTemplate          string                         `yaml:"analyze_fallback_prompt_template"`
	TitleCustomOutputRequirementsSection   string                         `yaml:"title_custom_output_requirements_section"`
	TitleOutputExample                     string                         `yaml:"title_output_example"`
	TitleFallbackPromptTemplate            string                         `yaml:"title_fallback_prompt_template"`
	OutlineOutputExample                   string                         `yaml:"outline_output_example"`
	OutlineCustomOutputRequirementsSection string                         `yaml:"outline_custom_output_requirements_section"`
	OutlineFallbackPromptTemplate          string                         `yaml:"outline_fallback_prompt_template"`
	OutlineResponseFormatDescription       string                         `yaml:"outline_response_format_description"`
	PromptRuleTaskNameSupplementFormat     string                         `yaml:"prompt_rule_task_name_supplement_format"`
	PromptRuleSceneSupplementFormat        string                         `yaml:"prompt_rule_scene_supplement_format"`
	PromptRuleToneSupplementFormat         string                         `yaml:"prompt_rule_tone_supplement_format"`
	PromptRuleWordCountSupplementFormat    string                         `yaml:"prompt_rule_word_count_supplement_format"`
	PromptRuleSupplementHeading            string                         `yaml:"prompt_rule_supplement_heading"`
	PromptRuleOutputRequirementsSection    string                         `yaml:"prompt_rule_output_requirements_section"`
	ArticleImitationPromptTemplate         string                         `yaml:"article_imitation_prompt_template"`
	ArticleImitationLocaleInstructions     map[string]string              `yaml:"article_imitation_locale_instructions"`
	ArticleImitationSettingLabels          map[string]string              `yaml:"article_imitation_setting_labels"`
	KnowledgePromptIntroLines              []string                       `yaml:"knowledge_prompt_intro_lines"`
	KnowledgeSnippetLabelFormat            string                         `yaml:"knowledge_snippet_label_format"`
	KnowledgeWebsiteMarkdown               knowledgeWebsiteMarkdownConfig `yaml:"knowledge_website_markdown"`
}

type knowledgeWebsiteMarkdownConfig struct {
	AssistantRole    string   `yaml:"assistant_role"`
	Task             string   `yaml:"task"`
	Requirements     []string `yaml:"requirements"`
	TitlePrefix      string   `yaml:"title_prefix"`
	URLPrefix        string   `yaml:"url_prefix"`
	ScopeSingle      string   `yaml:"scope_single"`
	ScopeMultiFormat string   `yaml:"scope_multi_format"`
}

type platformTemplatePromptEntry struct {
	PromptTemplate        string `yaml:"prompt_template"`
	AnalyzePromptTemplate string `yaml:"analyze_prompt_template"`
	TitlePromptTemplate   string `yaml:"title_prompt_template"`
	OutlinePromptTemplate string `yaml:"outline_prompt_template"`
}

type promptsCache struct {
	path    string
	modTime time.Time
	size    int64
	cfg     *promptsConfig
}

var (
	promptsMu          sync.RWMutex
	configuredFilePath string
	cachedPrompts      promptsCache
)

func SetConfigFile(path string) {
	promptsMu.Lock()
	defer promptsMu.Unlock()

	configuredFilePath = strings.TrimSpace(path)
	cachedPrompts = promptsCache{}
}

func loadPromptsConfig() *promptsConfig {
	path := resolvePromptsConfigFile()

	info, statErr := os.Stat(path)

	promptsMu.RLock()
	cached := cachedPrompts
	promptsMu.RUnlock()

	if statErr == nil && cached.cfg != nil && cached.path == path && cached.modTime.Equal(info.ModTime()) && cached.size == info.Size() {
		return cached.cfg
	}
	if statErr != nil && cached.cfg != nil && cached.path == path {
		return cached.cfg
	}

	promptsMu.Lock()
	defer promptsMu.Unlock()

	if info, statErr = os.Stat(path); statErr == nil {
		if cachedPrompts.cfg != nil && cachedPrompts.path == path && cachedPrompts.modTime.Equal(info.ModTime()) && cachedPrompts.size == info.Size() {
			return cachedPrompts.cfg
		}
	} else if cachedPrompts.cfg != nil && cachedPrompts.path == path {
		return cachedPrompts.cfg
	}

	if statErr != nil {
		panic(fmt.Sprintf("prompts: stat config %q failed: %v", path, statErr))
	}

	data, err := os.ReadFile(path)
	if err != nil {
		if cachedPrompts.cfg != nil && cachedPrompts.path == path {
			return cachedPrompts.cfg
		}
		panic(fmt.Sprintf("prompts: read config %q failed: %v", path, err))
	}

	var cfg promptsConfig
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		if cachedPrompts.cfg != nil && cachedPrompts.path == path {
			return cachedPrompts.cfg
		}
		panic(fmt.Sprintf("prompts: parse config %q failed: %v", path, err))
	}

	cachedPrompts = promptsCache{
		path:    path,
		modTime: info.ModTime(),
		size:    info.Size(),
		cfg:     &cfg,
	}

	return cachedPrompts.cfg
}

func loadRuntimePrompts() runtimePromptsConfig {
	return loadPromptsConfig().Runtime
}

func resolvePromptsConfigFile() string {
	promptsMu.RLock()
	configured := configuredFilePath
	promptsMu.RUnlock()

	if configured != "" {
		return configured
	}
	if value := strings.TrimSpace(os.Getenv(promptsConfigEnv)); value != "" {
		return value
	}
	if value := strings.TrimSpace(os.Getenv("CONFIG_PATH")); value != "" {
		return filepath.Join(filepath.Dir(value), "prompts.yml")
	}
	if _, file, _, ok := runtime.Caller(0); ok {
		return filepath.Clean(filepath.Join(filepath.Dir(file), "..", "..", "..", "configs", "prompts.yml"))
	}
	return filepath.Join("configs", "prompts.yml")
}

func platformTemplatePromptEntryFor(templateKey string) (platformTemplatePromptEntry, bool) {
	configs := loadPromptsConfig().PlatformTemplates
	if len(configs) == 0 {
		return platformTemplatePromptEntry{}, false
	}
	entry, ok := configs[strings.TrimSpace(templateKey)]
	return entry, ok
}

func ApplyPlatformTemplatePromptOverrides(templateKey string, promptTemplate *string, cardConfigJSON []byte) (*string, []byte) {
	entry, ok := platformTemplatePromptEntryFor(templateKey)
	if !ok {
		return clonePromptTemplate(promptTemplate), cloneBytes(cardConfigJSON)
	}

	nextPrompt := clonePromptTemplate(promptTemplate)
	if value := strings.TrimSpace(entry.PromptTemplate); value != "" {
		nextPrompt = &value
	}

	nextCardConfig := cloneBytes(cardConfigJSON)
	if len(nextCardConfig) == 0 {
		nextCardConfig = []byte("{}")
	}

	var payload map[string]interface{}
	if err := json.Unmarshal(nextCardConfig, &payload); err != nil {
		return nextPrompt, nextCardConfig
	}

	wizard, _ := payload["wizard"].(map[string]interface{})
	if wizard == nil {
		wizard = make(map[string]interface{})
		payload["wizard"] = wizard
	}

	if value := strings.TrimSpace(entry.AnalyzePromptTemplate); value != "" {
		wizard["analyze_prompt_template"] = value
	}
	if value := strings.TrimSpace(entry.TitlePromptTemplate); value != "" {
		wizard["title_prompt_template"] = value
	}
	if value := strings.TrimSpace(entry.OutlinePromptTemplate); value != "" {
		wizard["outline_prompt_template"] = value
	}

	marshaled, err := json.Marshal(payload)
	if err != nil {
		return nextPrompt, nextCardConfig
	}
	return nextPrompt, marshaled
}

func clonePromptTemplate(value *string) *string {
	if value == nil {
		return nil
	}
	cloned := *value
	return &cloned
}

func cloneBytes(value []byte) []byte {
	if len(value) == 0 {
		return nil
	}
	cloned := make([]byte, len(value))
	copy(cloned, value)
	return cloned
}