aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--README.md8
-rw-r--r--main.go213
2 files changed, 152 insertions, 69 deletions
diff --git a/README.md b/README.md
index 128ed6c..78e9bab 100644
--- a/README.md
+++ b/README.md
@@ -202,6 +202,14 @@ Determines which proxy to use to connect to the LLM endpoint:
llmProxy = "socks5://127.0.0.1:9050"
```
+### ircdName
+
+Name of the milla instance, must be unique across all instances.
+
+### adminOnly
+
+Milla will only answer if the nick is in the admin list.
+
### Example Config File
```toml
diff --git a/main.go b/main.go
index 37d87c8..c97b653 100644
--- a/main.go
+++ b/main.go
@@ -63,6 +63,7 @@ type TomlConfig struct {
DatabaseName string `toml:"databaseName"`
LLMProxy string `toml:"llmProxy"`
IRCProxy string `toml:"ircProxy"`
+ IRCDName string `toml:"ircdName"`
Temp float64 `toml:"temp"`
RequestTimeout int `toml:"requestTimeout"`
MillaReconnectDelay int `toml:"millaReconnectDelay"`
@@ -80,6 +81,7 @@ type TomlConfig struct {
AllowFlood bool `toml:"allowFlood"`
Debug bool `toml:"debug"`
Out bool `toml:"out"`
+ AdminOnly bool `toml:"adminOnly"`
Admins []string `toml:"admins"`
IrcChannels []string `toml:"ircChannels"`
ScrapeChannels []string `toml:"scrapeChannels"`
@@ -89,32 +91,73 @@ type AppConfig struct {
Ircd map[string]TomlConfig `toml:"ircd"`
}
-func NewTomlConfig() *TomlConfig {
- return &TomlConfig{
- IrcNick: "milla",
- IrcSaslUser: "milla",
- ChromaStyle: "rose-pine-moon",
- ChromaFormatter: "noop",
- Provider: "ollama",
- DatabaseAddress: "postgres",
- DatabaseUser: "milla",
- DatabaseName: "milladb",
- Temp: 0.5, //nolint:gomnd
- RequestTimeout: 10, //nolint:gomnd
- MillaReconnectDelay: 30, //nolint:gomnd
- IrcPort: 6697, //nolint:gomnd
- KeepAlive: 600, //nolint:gomnd
- MemoryLimit: 20, //nolint:gomnd
- PingDelay: 20, //nolint:gomnd
- PingTimeout: 20, //nolint:gomnd
- TopP: 0.9, //nolint:gomnd
- EnableSasl: false,
- SkipTLSVerify: false,
- UseTLS: true,
- AllowFlood: false,
- DisableSTSFallback: true,
- Debug: false,
- Out: false,
+func addSaneDefaults(config *TomlConfig) {
+ if config.IrcNick == "" {
+ config.IrcNick = "milla"
+ }
+
+ if config.IrcSaslUser == "" {
+ config.IrcSaslUser = "milla"
+ }
+
+ if config.ChromaStyle == "" {
+ config.ChromaStyle = "rose-pine-moon"
+ }
+
+ if config.ChromaFormatter == "" {
+ config.ChromaFormatter = "noop"
+ }
+
+ if config.Provider == "" {
+ config.Provider = "ollam"
+ }
+
+ if config.DatabaseAddress == "" {
+ config.DatabaseAddress = "postgres"
+ }
+
+ if config.DatabaseUser == "" {
+ config.DatabaseUser = "milla"
+ }
+
+ if config.DatabaseName == "" {
+ config.DatabaseName = "milladb"
+ }
+
+ if config.Temp == 0 {
+ config.Temp = 0.5 //nollint:gomnd
+ }
+
+ if config.RequestTimeout == 0 {
+ config.RequestTimeout = 10 //nolint:gomnd
+ }
+
+ if config.MillaReconnectDelay == 0 {
+ config.MillaReconnectDelay = 30 //nolint:gomnd
+ }
+
+ if config.IrcPort == 0 {
+ config.IrcPort = 6697 //nolint:gomnd
+ }
+
+ if config.KeepAlive == 0 {
+ config.KeepAlive = 600 //nolint:gomnd
+ }
+
+ if config.MemoryLimit == 0 {
+ config.MemoryLimit = 20 //nolint:gomnd
+ }
+
+ if config.PingDelay == 0 {
+ config.PingDelay = 20 //nolint:gomnd
+ }
+
+ if config.PingTimeout == 0 {
+ config.PingTimeout = 20 //nolint:gomnd
+ }
+
+ if config.TopP == 0. {
+ config.TopP = 0.9 //nolint:gomnd
}
}
@@ -144,6 +187,21 @@ type MemoryElement struct {
Content string `json:"content"`
}
+func getTableFromChanName(channel, ircdName string) string {
+ tableName := ircdName + "_" + channel
+ tableName = strings.ReplaceAll(tableName, "#", "")
+ tableName = strings.ReplaceAll(tableName, "-", "_")
+ tableName = strings.TrimSpace(tableName)
+
+ return tableName
+}
+
+func sanitizeLog(log string) string {
+ sanitizeLog := strings.ReplaceAll(log, "'", " ")
+
+ return sanitizeLog
+}
+
func returnGeminiResponse(resp *genai.GenerateContentResponse) string {
result := ""
@@ -365,6 +423,17 @@ func ollamaHandler(
if !strings.HasPrefix(event.Last(), appConfig.IrcNick+": ") {
return
}
+ if appConfig.AdminOnly {
+ byAdmin := false
+ for _, admin := range appConfig.Admins {
+ if event.Source.Name == admin {
+ byAdmin = true
+ }
+ }
+ if !byAdmin {
+ return
+ }
+ }
prompt := strings.TrimPrefix(event.Last(), appConfig.IrcNick+": ")
log.Println(prompt)
@@ -451,6 +520,8 @@ func ollamaHandler(
}
defer response.Body.Close()
+ log.Println("response body:", response.Body)
+
var writer bytes.Buffer
var ollamaChatResponse OllamaChatMessagesResponse
@@ -491,6 +562,17 @@ func geminiHandler(
if !strings.HasPrefix(event.Last(), appConfig.IrcNick+": ") {
return
}
+ if appConfig.AdminOnly {
+ byAdmin := false
+ for _, admin := range appConfig.Admins {
+ if event.Source.Name == admin {
+ byAdmin = true
+ }
+ }
+ if !byAdmin {
+ return
+ }
+ }
prompt := strings.TrimPrefix(event.Last(), appConfig.IrcNick+": ")
log.Println(prompt)
@@ -503,31 +585,6 @@ func geminiHandler(
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second)
defer cancel()
- // api and http client dont work together
- // https://github.com/google/generative-ai-go/issues/80
-
- // httpClient := http.Client{}
- // allProxy := os.Getenv("ALL_PROXY")
- // if allProxy != "" {
- // proxyUrl, err := url.Parse(allProxy)
- // if err != nil {
- // client.Cmd.ReplyTo(event, fmt.Sprintf("error: %s", err.Error()))
-
- // return
- // }
- // transport := &http.Transport{
- // Proxy: http.ProxyURL(proxyUrl),
- // }
-
- // httpClient.Transport = transport
- // }
-
- // clientGemini, err := genai.NewClient(ctx, option.WithAPIKey(appConfig.Apikey), option.WithHTTPClient(&httpClient))
-
- if appConfig.Apikey == "" {
- appConfig.Apikey = os.Getenv("MILLA_APIKEY")
- }
-
clientGemini, err := genai.NewClient(ctx, option.WithAPIKey(appConfig.Apikey))
if err != nil {
client.Cmd.ReplyTo(event, fmt.Sprintf("error: %s", err.Error()))
@@ -599,6 +656,17 @@ func chatGPTHandler(
if !strings.HasPrefix(event.Last(), appConfig.IrcNick+": ") {
return
}
+ if appConfig.AdminOnly {
+ byAdmin := false
+ for _, admin := range appConfig.Admins {
+ if event.Source.Name == admin {
+ byAdmin = true
+ }
+ }
+ if !byAdmin {
+ return
+ }
+ }
prompt := strings.TrimPrefix(event.Last(), appConfig.IrcNick+": ")
log.Println(prompt)
@@ -685,7 +753,7 @@ func chatGPTHandler(
})
}
-func connectToDB(appConfig TomlConfig, context *context.Context, poolChan chan *pgxpool.Pool) {
+func connectToDB(appConfig TomlConfig, ctx *context.Context, poolChan chan *pgxpool.Pool) {
for {
if appConfig.DatabaseUser == "" {
appConfig.DatabaseUser = os.Getenv("MILLA_DB_USER")
@@ -710,42 +778,48 @@ func connectToDB(appConfig TomlConfig, context *context.Context, poolChan chan *
appConfig.DatabaseAddress,
appConfig.DatabaseName)
- conn, err := pgxpool.New(*context, dbURL)
+ log.Println("dbURL:", dbURL)
+
+ poolConfig, err := pgxpool.ParseConfig(dbURL)
+ if err != nil {
+ log.Println(err)
+ }
+
+ pool, err := pgxpool.NewWithConfig(*ctx, poolConfig)
if err != nil {
log.Println(err)
time.Sleep(time.Duration(appConfig.MillaReconnectDelay) * time.Second)
} else {
- for _, channel := range appConfig.ScrapeChannels {
- query := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (id SERIAL PRIMARY KEY,channel TEXT NOT NULL,log TEXT NOT NULL,nick TEXT NOT NULL,dateadded TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
- strings.ReplaceAll(channel, "#", ""))
-
- log.Println(query)
+ log.Printf("%s connected to database", appConfig.IRCDName)
- _, err = conn.Query(*context, query)
+ for _, channel := range appConfig.ScrapeChannels {
+ tableName := getTableFromChanName(channel, appConfig.IRCDName)
+ query := fmt.Sprintf("create table if not exists %s (id serial primary key,channel text not null,log text not null,nick text not null,dateadded timestamp default current_timestamp)", tableName)
+ _, err = pool.Exec(*ctx, query)
if err != nil {
log.Println(err.Error())
time.Sleep(time.Duration(appConfig.MillaReconnectDelay) * time.Second)
}
}
- poolChan <- conn
+ poolChan <- pool
}
}
}
-func scrapeChannel(irc *girc.Client, poolChan chan *pgxpool.Pool) {
+func scrapeChannel(irc *girc.Client, poolChan chan *pgxpool.Pool, appConfig TomlConfig) {
irc.Handlers.AddBg(girc.PRIVMSG, func(client *girc.Client, event girc.Event) {
pool := <-poolChan
- query := fmt.Sprintf("INSERT INTO %s (channel,log,nick) VALUES ('%s','%s','%s')",
- strings.ReplaceAll(event.Params[0], "#", ""),
+ tableName := getTableFromChanName(event.Params[0], appConfig.IRCDName)
+ query := fmt.Sprintf(
+ "insert into %s (channel,log,nick) values ('%s','%s','%s')",
+ tableName,
event.Params[0],
event.Last(),
event.Source.Name,
)
- log.Println(query)
- _, err := pool.Query(
- context.Background(), query)
+ _, err := pool.Exec(context.Background(), query)
if err != nil {
log.Println(err.Error())
}
@@ -858,7 +932,7 @@ func runIRC(appConfig TomlConfig) {
}
})
- go scrapeChannel(irc, poolChan)
+ go scrapeChannel(irc, poolChan, appConfig)
}
for {
@@ -907,12 +981,13 @@ func main() {
}
for k, v := range config.Ircd {
+ addSaneDefaults(&v)
+ v.IRCDName = k
+ config.Ircd[k] = v
log.Println(k, v)
}
for _, v := range config.Ircd {
- log.Println(v)
-
go runIRC(v)
}