-
Notifications
You must be signed in to change notification settings - Fork 4
/
main.go
executable file
·124 lines (110 loc) · 3.11 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
package main
// main.go
import (
"encoding/json"
"fmt"
"github.com/joho/godotenv"
"github.com/meinside/openai-go"
log "github.com/sirupsen/logrus"
"github.com/tectiv3/chatgpt-bot/i18n"
"golang.org/x/crypto/ssh/terminal"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
stdlog "log"
"os"
"path"
"runtime"
"strings"
"time"
)
var (
Log *log.Entry
l *i18n.Localizer
// Version string will be set by linker
Version = "dev"
BuildTime = "unknown"
)
func main() {
logrus := log.New()
// redirect Go standard log library calls to logrus writer
stdlog.SetFlags(0)
stdlog.SetFlags(stdlog.LstdFlags | stdlog.Lshortfile)
logrus.Formatter = &log.TextFormatter{
FullTimestamp: true,
DisableTimestamp: !terminal.IsTerminal(int(os.Stdout.Fd())),
TimestampFormat: "Jan 2 15:04:05.000",
CallerPrettyfier: func(f *runtime.Frame) (string, string) {
return strings.Replace(f.Function, "(*Server).", "", -1),
fmt.Sprintf("%s:%d", path.Base(f.File), f.Line)
},
}
logrus.SetFormatter(logrus.Formatter)
logrus.SetReportCaller(true)
stdlog.SetOutput(logrus.Writer())
logrus.SetOutput(os.Stdout)
Log = logrus.WithFields(log.Fields{"ver": Version})
confFilepath := "config.json"
if len(os.Args) == 2 {
confFilepath = os.Args[1]
}
if conf, err := loadConfig(confFilepath); err == nil {
apiKey := conf.OpenAIAPIKey
orgID := conf.OpenAIOrganizationID
level := logger.Error
if conf.Verbose {
// level = logger.Info
logrus.SetLevel(log.DebugLevel)
}
newLogger := logger.New(
stdlog.New(os.Stdout, "\r\n", stdlog.LstdFlags), // io writer
logger.Config{
SlowThreshold: time.Second, // Slow SQL threshold
LogLevel: level, // Log level
IgnoreRecordNotFoundError: true, // Ignore ErrRecordNotFound error for logger
ParameterizedQueries: false, // Don't include params in the SQL log
Colorful: true, // Disable color
},
)
db, err := gorm.Open(sqlite.Open("bot.db"), &gorm.Config{Logger: newLogger})
if err != nil {
panic("failed to connect database")
}
// Migrate the schema
if err := db.AutoMigrate(&User{}); err != nil {
panic("failed to migrate user")
}
if err := db.AutoMigrate(&Chat{}); err != nil {
panic("failed to migrate chat")
}
if err := db.AutoMigrate(&ChatMessage{}); err != nil {
panic("failed to migrate chat message")
}
if err := db.AutoMigrate(&Role{}); err != nil {
panic("failed to migrate role")
}
Log.WithField("allowed_users", len(conf.AllowedTelegramUsers)).Info("Started")
server := &Server{
conf: conf,
ai: openai.NewClient(apiKey, orgID),
db: db,
}
l = i18n.New("ru", "en")
server.run()
} else {
Log.Warn("failed to load config", "error=", err)
}
}
// load config at given path
func loadConfig(fpath string) (conf config, err error) {
if err := godotenv.Load(); err != nil {
log.WithField("error", err).Warn("Error loading .env file")
}
var bytes []byte
if bytes, err = os.ReadFile(fpath); err == nil {
if err = json.Unmarshal(bytes, &conf); err == nil {
return conf, nil
}
}
return config{}, err
}