milla can run more than one instance of itself

4 files changed, 141 insertions, 48 deletions

diff --git a/README.md b/README.md
index a0337f0..128ed6c 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,8 @@
 # milla
 
 Milla is an IRC bot that sends things over to an LLM when you ask it questions and prints the answer with optional syntax-highlighting.<br/>
-Currently Supported:
+Milla can run more than one instance of itself, use different proxies(socks5 and http), connect to more than one IRC networks and log to different databases.<br/>
+Currently supported providers:
 
 - Ollama
 - Openai
@@ -19,7 +20,7 @@ Usage of ./milla:
 
 ## Config
 
-An exhaustive example is in `config-example.toml`.
+An example is provided under `config-example.toml`. Please note that all the config options are specific to one instance which is defined by `ircd.nameofyourinstance`.<br/>
 
 #### ircServer
 
@@ -201,6 +202,67 @@ Determines which proxy to use to connect to the LLM endpoint:
 llmProxy = "socks5://127.0.0.1:9050"
 ```
 
+### Example Config File
+
+```toml
+[ircd.devinet]
+ircServer = "irc.myawesomeircnet.com"
+ircPort = 6697
+ircNick = "milla"
+enableSasl = true
+ircSaslUser = "milla"
+ircSaslPass = "xxxxx"
+ircChannels = ["##chan1", "##chan2"]
+temp = 0.2
+requestTimeout = 10
+millaReconnectDelay = 60
+model = "gpt-3.5-turbo"
+chromaStyle = "rose-pine-moon"
+chromaFormatter = "terminal256"
+provider = "chatgpt"
+apikey = "xxxx"
+memoryLimit = 20
+admins = ["noone_has_this_nick"]
+debug = true
+out = true
+databaseAddress = "postgres:5432"
+databasePassword = "changeme"
+databaseUser = "devi"
+databaseName = "milla"
+scrapeChannels = ["#soulhack", "#warroom"]
+ircProxy = "socks5://127.0.0.1:9050"
+llmProxy = "http://127.0.0.1:8180"
+skipTLSVerify = false
+useTLS = true
+
+[ircd.liberanet]
+ircServer = "irc.libera.chat"
+ircNick = "milla"
+model = "gpt-3.5-turbo"
+ircPort = 6697
+chromaStyle = "rose-pine-moon"
+chromaFormatter = "terminal16m"
+provider = "gemini"
+apikey = "xxxx"
+temp = 0.5
+requestTimeout = 10
+millaReconnectDelay = 60
+keepAlive = 20
+memoryLimit = 20
+pingDelay = 20
+pingTimeout = 600
+skipTLSVerify = false
+useTLS = true
+disableSTSFallback = true
+allowFlood = false
+admins = ["noone_has_this_nick"]
+ircChannels = ["##milla1", "##milla2"]
+debug = true
+out = true
+ircProxy = "socks5://127.0.0.1:9051"
+llmProxy = "http://127.0.0.1:8181"
+```
+
 ## Commands
 
 #### help
diff --git a/config-example.toml b/config-example.toml
index 9a5a56e..60f32a4 100644
--- a/config-example.toml
+++ b/config-example.toml
@@ -1,33 +1,56 @@
-ircServer = "irc.libera.caht"
+[ircd.devinet]
+ircServer = "irc.myawesomeircnet.com"
+ircPort = 6697
 ircNick = "milla"
-ircSaslUser = ""
-ircSaslPass = ""
-ollamaEndpoint = "http://127.0.0.1:11434/api/chat"
-model = "ollama2-uncensored"
-chromaStyle = "monokai"
+enableSasl = true
+ircSaslUser = "milla"
+ircSaslPass = "xxxxx"
+ircChannels = ["##chan1", "##chan2"]
+temp = 0.2
+requestTimeout = 10
+millaReconnectDelay = 60
+model = "gpt-3.5-turbo"
+chromaStyle = "rose-pine-moon"
+chromaFormatter = "terminal256"
+provider = "chatgpt"
+apikey = "xxxx"
+memoryLimit = 20
+admins = ["noone_has_this_nick"]
+debug = true
+out = true
+databaseAddress = "postgres:5432"
+databasePassword = "changeme"
+databaseUser = "devi"
+databaseName = "milla"
+scrapeChannels = ["#soulhack", "#warroom"]
+ircProxy = "socks5://127.0.0.1:9050"
+llmProxy = "http://127.0.0.1:8180"
+skipTLSVerify = false
+useTLS = true
+
+[ircd.liberanet]
+ircServer = "irc.libera.chat"
+ircNick = "milla"
+model = "gpt-3.5-turbo"
+ircPort = 6697
+chromaStyle = "rose-pine-moon"
 chromaFormatter = "terminal16m"
-provider = "ollama"
-apikey = "xxx"
-ollamaSystem = ""
-clientCertPath = ""
-serverPass = ""
-bind = ""
+provider = "gemini"
+apikey = "xxxx"
 temp = 0.5
 requestTimeout = 10
 millaReconnectDelay = 60
-ircPort = 6669
 keepAlive = 20
 memoryLimit = 20
 pingDelay = 20
 pingTimeout = 600
-topP = 0.9
-topK = 0
-enableSasl = false
 skipTLSVerify = false
 useTLS = true
 disableSTSFallback = true
 allowFlood = false
-debug = false
-out = false
-admins = ["nick1", "nick2"]
+admins = ["noone_has_this_nick"]
 ircChannels = ["##milla1", "##milla2"]
+debug = true
+out = true
+ircProxy = "socks5://127.0.0.1:9051"
+llmProxy = "http://127.0.0.1:8181"
diff --git a/docker-compose-postgres.yaml b/docker-compose-postgres.yaml
index a10e79f..38a0924 100644
--- a/docker-compose-postgres.yaml
+++ b/docker-compose-postgres.yaml
@@ -19,7 +19,7 @@ services:
     entrypoint: ["/usr/bin/milla"]
     command: ["--config", "/config.toml"]
     volumes:
-      - ./config-gpt.toml:/config.toml
+      - ./config.toml:/config.toml
       - /etc/localtime:/etc/localtime:ro
     cap_drop:
       - ALL
diff --git a/main.go b/main.go
index e2e5e93..37d87c8 100644
--- a/main.go
+++ b/main.go
@@ -13,11 +13,13 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"os/signal"
 	"reflect"
 	"regexp"
 	"runtime"
 	"strconv"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/BurntSushi/toml"
@@ -37,7 +39,6 @@ var (
 	errCantSet           = errors.New("can't set field")
 	errWrongDataForField = errors.New("wrong data type for field")
 	errUnsupportedType   = errors.New("unsupported type")
-	dbConnection         *pgxpool.Pool //nolint:gochecknoglobals
 )
 
 type TomlConfig struct {
@@ -84,6 +85,10 @@ type TomlConfig struct {
 	ScrapeChannels      []string `toml:"scrapeChannels"`
 }
 
+type AppConfig struct {
+	Ircd map[string]TomlConfig `toml:"ircd"`
+}
+
 func NewTomlConfig() *TomlConfig {
 	return &TomlConfig{
 		IrcNick:             "milla",
@@ -680,7 +685,7 @@ func chatGPTHandler(
 	})
 }
 
-func connectToDB(appConfig TomlConfig, context *context.Context) {
+func connectToDB(appConfig TomlConfig, context *context.Context, poolChan chan *pgxpool.Pool) {
 	for {
 		if appConfig.DatabaseUser == "" {
 			appConfig.DatabaseUser = os.Getenv("MILLA_DB_USER")
@@ -723,18 +728,14 @@ func connectToDB(appConfig TomlConfig, context *context.Context) {
 				}
 			}
 
-			dbConnection = conn
+			poolChan <- conn
 		}
 	}
 }
 
-func scrapeChannel(irc *girc.Client) {
+func scrapeChannel(irc *girc.Client, poolChan chan *pgxpool.Pool) {
 	irc.Handlers.AddBg(girc.PRIVMSG, func(client *girc.Client, event girc.Event) {
-		if dbConnection == nil {
-			log.Println("missed logging message because currently not connected to db")
-
-			return
-		}
+		pool := <-poolChan
 		query := fmt.Sprintf("INSERT INTO %s (channel,log,nick) VALUES ('%s','%s','%s')",
 			strings.ReplaceAll(event.Params[0], "#", ""),
 			event.Params[0],
@@ -743,7 +744,7 @@ func scrapeChannel(irc *girc.Client) {
 		)
 		log.Println(query)
 
-		_, err := dbConnection.Query(
+		_, err := pool.Query(
 			context.Background(), query)
 		if err != nil {
 			log.Println(err.Error())
@@ -751,13 +752,15 @@ func scrapeChannel(irc *girc.Client) {
 	})
 }
 
-func runIRC(appConfig TomlConfig, ircChan chan *girc.Client, dbChan chan *pgxpool.Pool) {
+func runIRC(appConfig TomlConfig) {
 	var OllamaMemory []MemoryElement
 
 	var GeminiMemory []*genai.Content
 
 	var GPTMemory []openai.ChatCompletionMessage
 
+	poolChan := make(chan *pgxpool.Pool, 1)
+
 	irc := girc.New(girc.Config{
 		Server:             appConfig.IrcServer,
 		Port:               appConfig.IrcPort,
@@ -841,10 +844,12 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client, dbChan chan *pgxpoo
 		chatGPTHandler(irc, &appConfig, &GPTMemory)
 	}
 
-	context, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second)
-	defer cancel()
+	if appConfig.DatabaseAddress != "" {
+		context, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second)
+		defer cancel()
 
-	go connectToDB(appConfig, &context)
+		go connectToDB(appConfig, &context, poolChan)
+	}
 
 	if len(appConfig.ScrapeChannels) > 0 {
 		irc.Handlers.AddBg(girc.CONNECTED, func(c *girc.Client, e girc.Event) {
@@ -853,9 +858,8 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client, dbChan chan *pgxpoo
 			}
 		})
 
-		go scrapeChannel(irc)
+		go scrapeChannel(irc, poolChan)
 	}
-	ircChan <- irc
 
 	for {
 		var dialer proxy.Dialer
@@ -863,15 +867,11 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client, dbChan chan *pgxpoo
 		if appConfig.IRCProxy != "" {
 			proxyURL, err := url.Parse(appConfig.IRCProxy)
 			if err != nil {
-				cancel()
-
 				log.Fatal(err.Error())
 			}
 
 			dialer, err = proxy.FromURL(proxyURL, &net.Dialer{Timeout: time.Duration(appConfig.RequestTimeout) * time.Second})
 			if err != nil {
-				cancel()
-
 				log.Fatal(err.Error())
 			}
 		}
@@ -887,6 +887,9 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client, dbChan chan *pgxpoo
 }
 
 func main() {
+	quitChannel := make(chan os.Signal, 1)
+	signal.Notify(quitChannel, syscall.SIGINT, syscall.SIGTERM)
+
 	configPath := flag.String("config", "./config.toml", "path to the config file")
 
 	flag.Parse()
@@ -896,17 +899,22 @@ func main() {
 		log.Fatal(err)
 	}
 
-	appConfig := NewTomlConfig()
+	var config AppConfig
 
-	_, err = toml.Decode(string(data), &appConfig)
+	_, err = toml.Decode(string(data), &config)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	log.Println(appConfig)
+	for k, v := range config.Ircd {
+		log.Println(k, v)
+	}
+
+	for _, v := range config.Ircd {
+		log.Println(v)
 
-	ircChan := make(chan *girc.Client, 1)
-	dbConn := make(chan *pgxpool.Pool, 1)
+		go runIRC(v)
+	}
 
-	runIRC(*appConfig, ircChan, dbConn)
+	<-quitChannel
 }