347 lines
9.3 KiB
Go
347 lines
9.3 KiB
Go
package main
|
||
|
||
import (
|
||
"bufio"
|
||
"embed"
|
||
"fmt"
|
||
"io/ioutil"
|
||
"os"
|
||
"strconv"
|
||
"strings"
|
||
|
||
"github.com/voocel/ainovel-cli/app"
|
||
"github.com/voocel/ainovel-cli/tools"
|
||
"github.com/voocel/ainovel-cli/tui"
|
||
)
|
||
|
||
//go:embed prompts/*.md
|
||
var promptsFS embed.FS
|
||
|
||
//go:embed references
|
||
var referencesFS embed.FS
|
||
|
||
//go:embed styles/*.md
|
||
var stylesFS embed.FS
|
||
|
||
func main() {
|
||
// 处理 start 子命令
|
||
if len(os.Args) > 1 && os.Args[1] == "start" {
|
||
handleStartCommand()
|
||
return
|
||
}
|
||
|
||
style := envOr("NOVEL_STYLE", "default")
|
||
refs := loadReferences(style)
|
||
prompts := loadPrompts()
|
||
styles := loadStyles()
|
||
cfg := buildConfig(style)
|
||
|
||
prompt := parsePrompt()
|
||
if prompt != "" {
|
||
// CLI 模式：有命令行参数，直接运行
|
||
cfg.Prompt = prompt
|
||
if err := app.Run(cfg, refs, prompts, styles); err != nil {
|
||
fmt.Fprintf(os.Stderr, "error: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
return
|
||
}
|
||
|
||
// TUI 模式：无命令行参数，启动交互界面
|
||
if err := tui.Run(cfg, refs, prompts, styles); err != nil {
|
||
fmt.Fprintf(os.Stderr, "error: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
}
|
||
|
||
func buildConfig(style string) app.Config {
|
||
provider := envOr("LLM_PROVIDER", "openrouter")
|
||
apiKey := os.Getenv("Z_OPENAI_API_KEY")
|
||
baseURL := os.Getenv("Z_OPENAI_BASE_URL")
|
||
switch provider {
|
||
case "anthropic":
|
||
apiKey = envOr("ANTHROPIC_API_KEY", apiKey)
|
||
baseURL = envOr("ANTHROPIC_BASE_URL", baseURL)
|
||
case "gemini":
|
||
apiKey = envOr("GEMINI_API_KEY", apiKey)
|
||
baseURL = envOr("GEMINI_BASE_URL", baseURL)
|
||
case "openrouter":
|
||
apiKey = envOr("OPENROUTER_API_KEY", apiKey)
|
||
baseURL = envOr("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
|
||
}
|
||
|
||
cfg := app.Config{
|
||
NovelName: envOr("NOVEL_NAME", ""),
|
||
Provider: provider,
|
||
APIKey: apiKey,
|
||
BaseURL: baseURL,
|
||
ModelName: envOr("LLM_MODEL", "stepfun/step-3.5-flash:free"),
|
||
Style: style,
|
||
}
|
||
return cfg
|
||
}
|
||
|
||
func parsePrompt() string {
|
||
if len(os.Args) < 2 {
|
||
return ""
|
||
}
|
||
if os.Args[1] == "-help" || os.Args[1] == "--help" || os.Args[1] == "-h" {
|
||
printHelp()
|
||
os.Exit(0)
|
||
}
|
||
return strings.Join(os.Args[1:], " ")
|
||
}
|
||
|
||
func printHelp() {
|
||
fmt.Println(`ainovel-cli - AI 小说生成工具
|
||
|
||
用法:
|
||
ainovel-cli start 交互式开始：选择小说 → 选择风格 → 启动 TUI
|
||
ainovel-cli [prompt] CLI 模式：直接生成小说
|
||
ainovel-cli TUI 模式：启动交互界面
|
||
|
||
环境变量:
|
||
NOVEL_NAME 小说名称（默认：novel）
|
||
NOVEL_STYLE 小说风格（默认：default）
|
||
可选值：
|
||
default 通用风格，叙事张弛有度，五感描写，对话自然
|
||
fantasy 奇幻冒险，世界观自然展开，魔法体系有代价感
|
||
romance 言情，情感递进有节奏，关系张力与内心描写并重
|
||
suspense 悬疑推理，多线叙事，信息差悬念，线索管理严谨
|
||
LLM_PROVIDER LLM 提供商：openrouter|anthropic|gemini（默认：openrouter）
|
||
LLM_MODEL 模型名称（默认：stepfun/step-3.5-flash:free）
|
||
|
||
Z_OPENAI_API_KEY API 密钥
|
||
Z_OPENAI_BASE_URL API 地址
|
||
|
||
ANTHROPIC_API_KEY Anthropic API 密钥
|
||
ANTHROPIC_BASE_URL Anthropic API 地址
|
||
|
||
GEMINI_API_KEY Gemini API 密钥
|
||
GEMINI_BASE_URL Gemini API 地址
|
||
|
||
OPENROUTER_API_KEY OpenRouter API 密钥
|
||
OPENROUTER_BASE_URL OpenRouter API 地址
|
||
|
||
示例:
|
||
ainovel-cli "写一部科幻小说，主角是时间旅行者"
|
||
NOVEL_NAME=时空之旅 ainovel-cli "写一部科幻小说"
|
||
ainovel-cli # 启动 TUI 交互模式
|
||
|
||
tmux 使用（推荐后台运行）:
|
||
tmux new -s novel # 新建名为 novel 的会话
|
||
ainovel-cli "写一部科幻小说" # 在会话中运行，关闭终端不中断
|
||
Ctrl+b d # 脱离会话（程序继续跑）
|
||
tmux attach -t novel # 随时回来查看进度
|
||
tmux ls # 列出所有会话
|
||
tmux kill-session -t novel # 结束会话`)
|
||
}
|
||
|
||
func loadReferences(style string) tools.References {
|
||
refs := tools.References{
|
||
ChapterGuide: mustRead(referencesFS, "references/chapter-guide.md"),
|
||
HookTechniques: mustRead(referencesFS, "references/hook-techniques.md"),
|
||
QualityChecklist: mustRead(referencesFS, "references/quality-checklist.md"),
|
||
OutlineTemplate: mustRead(referencesFS, "references/outline-template.md"),
|
||
CharacterTemplate: mustRead(referencesFS, "references/character-template.md"),
|
||
ChapterTemplate: mustRead(referencesFS, "references/chapter-template.md"),
|
||
Consistency: mustRead(referencesFS, "references/consistency.md"),
|
||
ContentExpansion: mustRead(referencesFS, "references/content-expansion.md"),
|
||
DialogueWriting: mustRead(referencesFS, "references/dialogue-writing.md"),
|
||
LongformPlanning: mustRead(referencesFS, "references/longform-planning.md"),
|
||
Differentiation: mustRead(referencesFS, "references/differentiation.md"),
|
||
}
|
||
if style != "" && style != "default" {
|
||
path := "references/" + style + "/style-references.md"
|
||
if data, err := referencesFS.ReadFile(path); err == nil {
|
||
refs.StyleReference = string(data)
|
||
}
|
||
}
|
||
return refs
|
||
}
|
||
|
||
func loadPrompts() app.Prompts {
|
||
return app.Prompts{
|
||
Coordinator: mustRead(promptsFS, "prompts/coordinator.md"),
|
||
ArchitectShort: mustRead(promptsFS, "prompts/architect-short.md"),
|
||
ArchitectMid: mustRead(promptsFS, "prompts/architect-mid.md"),
|
||
ArchitectLong: mustRead(promptsFS, "prompts/architect-long.md"),
|
||
Writer: mustRead(promptsFS, "prompts/writer.md"),
|
||
Editor: mustRead(promptsFS, "prompts/editor.md"),
|
||
}
|
||
}
|
||
|
||
func mustRead(fs embed.FS, path string) string {
|
||
data, err := fs.ReadFile(path)
|
||
if err != nil {
|
||
panic(fmt.Sprintf("embed read %s: %v", path, err))
|
||
}
|
||
return string(data)
|
||
}
|
||
|
||
func loadStyles() map[string]string {
|
||
styles := make(map[string]string)
|
||
entries, err := stylesFS.ReadDir("styles")
|
||
if err != nil {
|
||
return styles
|
||
}
|
||
for _, e := range entries {
|
||
if e.IsDir() {
|
||
continue
|
||
}
|
||
name := strings.TrimSuffix(e.Name(), ".md")
|
||
data, err := stylesFS.ReadFile("styles/" + e.Name())
|
||
if err != nil {
|
||
continue
|
||
}
|
||
styles[name] = string(data)
|
||
}
|
||
return styles
|
||
}
|
||
|
||
func envOr(key, fallback string) string {
|
||
if v := os.Getenv(key); v != "" {
|
||
return v
|
||
}
|
||
return fallback
|
||
}
|
||
|
||
// handleStartCommand 处理 start 子命令：交互式选择小说和风格
|
||
func handleStartCommand() {
|
||
// Step 1: 选择小说
|
||
novelName := selectNovel()
|
||
if novelName == "" {
|
||
fmt.Println("已取消")
|
||
os.Exit(0)
|
||
}
|
||
|
||
// Step 2: 选择风格
|
||
style := selectStyle()
|
||
if style == "" {
|
||
fmt.Println("已取消")
|
||
os.Exit(0)
|
||
}
|
||
|
||
// 设置环境变量并启动 TUI
|
||
os.Setenv("NOVEL_NAME", novelName)
|
||
os.Setenv("NOVEL_STYLE", style)
|
||
|
||
refs := loadReferences(style)
|
||
prompts := loadPrompts()
|
||
styles := loadStyles()
|
||
cfg := buildConfig(style)
|
||
cfg.NovelName = novelName
|
||
|
||
if err := tui.Run(cfg, refs, prompts, styles); err != nil {
|
||
fmt.Fprintf(os.Stderr, "error: %v\n", err)
|
||
os.Exit(1)
|
||
}
|
||
}
|
||
|
||
// selectNovel 显示小说选择菜单
|
||
func selectNovel() string {
|
||
options := []string{"📝 创建新小说"}
|
||
|
||
// 扫描 output 目录下现有小说
|
||
outputDir := "output"
|
||
entries, err := ioutil.ReadDir(outputDir)
|
||
novels := []string{}
|
||
if err == nil {
|
||
for _, e := range entries {
|
||
if e.IsDir() && e.Name() != "." && e.Name() != ".." {
|
||
novels = append(novels, e.Name())
|
||
options = append(options, "📖 "+e.Name())
|
||
}
|
||
}
|
||
}
|
||
|
||
if len(options) == 1 {
|
||
// 只有创建选项，直接输入新名称
|
||
return promptNovelName()
|
||
}
|
||
|
||
// 显示菜单
|
||
fmt.Println("\n=== 选择小说或创建新的 ===")
|
||
for i, opt := range options {
|
||
fmt.Printf("%d. %s\n", i+1, opt)
|
||
}
|
||
|
||
choice := promptChoice(len(options))
|
||
if choice < 0 {
|
||
return ""
|
||
}
|
||
|
||
if choice == 0 {
|
||
// 创建新小说
|
||
return promptNovelName()
|
||
}
|
||
|
||
// 返回现有小说名称
|
||
return novels[choice-1]
|
||
}
|
||
|
||
// selectStyle 显示风格选择菜单
|
||
func selectStyle() string {
|
||
styles := []struct {
|
||
name string
|
||
desc string
|
||
}{
|
||
{"default", "通用风格，叙事张弛有度，五感描写，对话自然"},
|
||
{"fantasy", "奇幻冒险，世界观自然展开，魔法体系有代价感"},
|
||
{"romance", "言情，情感递进有节奏，关系张力与内心描写并重"},
|
||
{"suspense", "悬疑推理，多线叙事，信息差悬念，线索管理严谨"},
|
||
}
|
||
|
||
fmt.Println("\n=== 选择小说风格 ===")
|
||
for i, s := range styles {
|
||
fmt.Printf("%d. %-10s %s\n", i+1, s.name, s.desc)
|
||
}
|
||
|
||
choice := promptChoice(len(styles))
|
||
if choice < 0 {
|
||
// EOF 或错误，使用默认风格
|
||
fmt.Println("使用默认风格...")
|
||
return "default"
|
||
}
|
||
return styles[choice].name
|
||
}
|
||
|
||
// promptNovelName 提示输入新小说名称
|
||
func promptNovelName() string {
|
||
reader := bufio.NewReader(os.Stdin)
|
||
fmt.Print("\n请输入小说名称 (默认: novel): ")
|
||
name, err := reader.ReadString('\n')
|
||
if err != nil {
|
||
return "novel"
|
||
}
|
||
name = strings.TrimSpace(name)
|
||
if name == "" {
|
||
return "novel"
|
||
}
|
||
return name
|
||
}
|
||
|
||
// promptChoice 提示用户选择菜单项
|
||
func promptChoice(maxChoice int) int {
|
||
reader := bufio.NewReader(os.Stdin)
|
||
for {
|
||
fmt.Print("\n请选择 (输入数字): ")
|
||
input, err := reader.ReadString('\n')
|
||
if err != nil {
|
||
// EOF 或其他错误，返回 -1 表示取消
|
||
return -1
|
||
}
|
||
input = strings.TrimSpace(input)
|
||
if input == "" {
|
||
continue
|
||
}
|
||
|
||
choice, err := strconv.Atoi(input)
|
||
if err != nil || choice < 1 || choice > maxChoice {
|
||
fmt.Printf("❌ 请输入 1 到 %d 之间的数字\n", maxChoice)
|
||
continue
|
||
}
|
||
return choice - 1 // 转换为 0-based 索引
|
||
}
|
||
}
|