72 lines
1.4 KiB
Go
72 lines
1.4 KiB
Go
|
package ai
|
||
|
|
|
||
|
|
import (
|
||
|
|
"embed"
|
||
|
|
"fmt"
|
||
|
|
"strings"
|
||
|
|
"text/template"
|
||
|
|
)
|
||
|
|
|
||
|
|
//go:embed prompts/*.txt
|
||
|
|
var promptFS embed.FS
|
||
|
|
|
||
|
|
// PromptManager 提示词管理器
|
||
|
|
type PromptManager struct {
|
||
|
|
templates map[string]*template.Template
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewPromptManager 创建提示词管理器
|
||
|
|
func NewPromptManager() *PromptManager {
|
||
|
|
pm := &PromptManager{
|
||
|
|
templates: make(map[string]*template.Template),
|
||
|
|
}
|
||
|
|
|
||
|
|
pm.loadTemplates()
|
||
|
|
return pm
|
||
|
|
}
|
||
|
|
|
||
|
|
func (pm *PromptManager) loadTemplates() {
|
||
|
|
templateFiles := map[string]string{
|
||
|
|
"mermaid": "prompts/mermaid.txt",
|
||
|
|
"guodegang": "prompts/guodegang.txt",
|
||
|
|
"quiz": "prompts/quiz.txt",
|
||
|
|
"card": "prompts/card.txt",
|
||
|
|
"scenario": "prompts/scenario.txt",
|
||
|
|
}
|
||
|
|
|
||
|
|
for name, path := range templateFiles {
|
||
|
|
content, err := promptFS.ReadFile(path)
|
||
|
|
if err != nil {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
|
||
|
|
tmpl, err := template.New(name).Parse(string(content))
|
||
|
|
if err != nil {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
|
||
|
|
pm.templates[name] = tmpl
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// GeneratePrompt 生成提示词
|
||
|
|
func (pm *PromptManager) GeneratePrompt(templateType string, knowledge string) (string, error) {
|
||
|
|
tmpl, ok := pm.templates[templateType]
|
||
|
|
if !ok {
|
||
|
|
return "", fmt.Errorf("template type %s not found", templateType)
|
||
|
|
}
|
||
|
|
|
||
|
|
var buf strings.Builder
|
||
|
|
data := struct {
|
||
|
|
Knowledge string
|
||
|
|
}{
|
||
|
|
Knowledge: knowledge,
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := tmpl.Execute(&buf, data); err != nil {
|
||
|
|
return "", fmt.Errorf("failed to execute template: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
return buf.String(), nil
|
||
|
|
}
|