ai-bot/main.go

package main

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/peterh/liner"

	"myaibot/internal/bot"
	"myaibot/internal/cli"
	"myaibot/internal/config"
	"myaibot/internal/store"
)

func main() {
	os.Setenv("TIKTOKEN_CACHE_DIR", filepath.Join("data", "tiktoken"))
	cfg, err := config.Load()
	if err != nil {
		log.Fatalf("加载配置失败: %v", err)
	}
	autoFetchModels(cfg)

	db, err := store.Open(&cfg.Database)
	if err != nil {
		log.Fatalf("数据库连接失败: %v", err)
	}
	fmt.Printf("💾 数据库已连接 (%s)\n", cfg.Database.Driver)

	b, err := bot.New(cfg, db)
	if err != nil {
		fmt.Printf("⚠️  %v\n", err)
		fmt.Println("请填写工具配置文件后重新启动。")
		store.Close(db)
		os.Exit(1)
	}

	if err := store.Migrate(db, cfg.Database.Driver); err != nil {
		log.Fatalf("数据库迁移失败: %v", err)
	}
	if sess, err := store.LoadLatestSession(db); err != nil {
		fmt.Printf("⚠️ 读取上次会话失败: %v\n", err)
	} else if sess != nil && len(sess.Messages) > 0 {
		b.RestoreSession(sess)
		fmt.Printf("💬 已恢复上次会话 (%d 条消息)\n", len(sess.Messages))
	}

	provider, model := b.Current()
	fmt.Printf("🤖 %s 已启动 (供应商: %s, 模型: %s)。输入问题开始对话，输入 /help 查看命令。\n",
		cfg.BotName, provider, model)

	line := liner.NewLiner()
	defer line.Close()
	line.SetCtrlCAborts(true)
	line.SetCompleter(func(s string) []string {
		return cli.Complete(s, b.Models())
	})

	h := cli.New(b, db)
	for {
		input, err := line.Prompt("你: ")
		if errors.Is(err, io.EOF) || errors.Is(err, liner.ErrPromptAborted) {
			fmt.Println("再见！")
			break
		}
		input = strings.TrimSpace(input)
		if input == "" {
			continue
		}
		line.AppendHistory(input)
		if strings.HasPrefix(input, "/") {
			if !h.Handle(input) {
				break
			}
			continue
		}
		fmt.Printf("%s: ", cfg.BotName)
		thinkStyle, resetStyle := false, false
		_, err = b.Chat(context.Background(), input,
			func(text string) {
				if !thinkStyle {
					fmt.Print("\x1b[3;90m🧠 ")
					thinkStyle, resetStyle = true, true
				}
				fmt.Print(text)
			},
			func(text string) {
				if resetStyle {
					fmt.Print("\x1b[0m\n")
					thinkStyle, resetStyle = false, false
				}
				fmt.Print(text)
			},
			func(name, args string) {
				if resetStyle {
					fmt.Print("\x1b[0m\n")
					thinkStyle, resetStyle = false, false
				}
				fmt.Printf("🔧 调用工具: %s %s\n", name, args)
			},
			func(text string) {
				if !thinkStyle {
					fmt.Print("\x1b[3;90m🔧 ")
					thinkStyle, resetStyle = true, true
				}
				fmt.Print(text)
			},
		)
		fmt.Println()
		if err != nil {
			fmt.Printf("⚠️  %v\n", err)
			continue
		}
	}

	saveSession(db, b)
	store.Close(db)
}

func autoFetchModels(cfg *config.Config) {
	for i := range cfg.Providers {
		p := &cfg.Providers[i]
		if !p.AutoFetchModels {
			continue
		}
		if p.APIKey == "" {
			fmt.Printf("⚠️  供应商 %s 启用了 auto_fetch_models 但未配置 api_key，跳过自动获取\n", p.Name)
			continue
		}
		before := append([]config.ModelConfig(nil), p.Models...)
		if err := config.FetchModels(context.Background(), p); err != nil {
			fmt.Printf("⚠️  自动获取模型失败 (供应商 %s): %v，使用现有模型列表\n", p.Name, err)
			continue
		}
		fmt.Printf("📚 已从 API 获取模型列表 (供应商 %s): %d 个模型\n", p.Name, len(p.Models))
		if !config.ModelsEqual(before, p.Models) {
			if err := config.Save(cfg); err != nil {
				fmt.Printf("⚠️  模型列表写回配置失败: %v\n", err)
			}
		}
	}
}

func saveSession(db *sql.DB, b *bot.Bot) {
	msgs := b.SessionMessages()
	if len(msgs) == 0 {
		return
	}
	provider, model := b.Current()
	sess := &store.Session{
		Provider:     provider,
		Model:        model,
		SystemPrompt: "",
		Messages:     msgs,
	}
	if _, err := store.SaveSession(db, sess); err != nil {
		fmt.Printf("⚠️  保存会话失败: %v\n", err)
		return
	}
	fmt.Printf("💾 会话已保存 (%d 条消息)\n", len(msgs))
}