fixes #24, fixes #25, fixes #29

2 files changed, 144 insertions, 105 deletions

diff --git a/README.md b/README.md
index fc49871..5ba8ed4 100644
--- a/README.md
+++ b/README.md
@@ -1,12 +1,17 @@
 # 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/>
-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:
+Milla is an IRC bot that:
 
-- Ollama
-- Openai
-- Gemini
+- sends things over to an LLM when you ask it questions and prints the answer with optional syntax-highlighting.<br/>
+  Currently supported providers:
+
+* Ollama
+* Openai
+* Gemini
+
+- Milla can run more than one instance of itself
+- Each instance can connect to a different ircd, and will get the full set of configs, e.g. different proxies, different postgres instance, ...
+- You can define custom commands in the form of SQL queries to the database with the SQL query result being passed to the bot along with the given prompt and an optional limit so you don't go bankrupt(unless you are running ollama locally like the smart cookie that you are).<br/>
 
 ![milla](./milla.png)
 
@@ -159,7 +164,7 @@ Whether to write raw messages to stdout.
 
 #### admins
 
-List of adimns for the bot. Only admins can use commands.
+List of admins for the bot. Only admins can use commands.
 
 ```
 admins = ["admin1", "admin2"]
@@ -173,25 +178,27 @@ List of channels for the bot to join when it connects to the server.
 ircChannels = ["#channel1", "#channel2"]
 ```
 
+Please note that the bot does not have to join a channel to be usable. One can simply query the bot directly as well.<br/>
+
 ### databaseUser
 
-Name of the database user. Can also be passed an an environment variable.
+Name of the database user.
 
 ### databasePassword
 
-Password for the database user. Can also be passed an an environment variable.
+Password for the database user.
 
 ### databaseAddress
 
-Address of the database. Can also be passed as and environment variable.
+Address of the database.
 
 ### databaseName
 
-Name of the database. Can also be passed as and environment variable.
+Name of the database.
 
 ### ircProxy
 
-Determines which proxy to use to connect to the irc network:
+Determines which proxy to use to connect to the IRC network:
 
 ```
 ircProxy = "socks5://127.0.0.1:9050"
@@ -230,7 +237,9 @@ webirc password to use.
 webirc address to use.
 
 ## Custom Commands
+
 Custom commands let you define a command that does a SQL query to the database and performs the given task. Here's an example:
+
 ```toml
 [ircd.devinet_terra.customCommands.digest]
 sql = "select log from liberanet_milla_us_market_news;"
@@ -241,10 +250,11 @@ sql= "select log from liberanet_milla_us_market_news;"
 limit= 300
 prompt= "given all the data, summarize the news for me"
 ```
+
 In the above example digest and summarize will be the names of the commands: `milla: /cmd summarize`.<br/>
 Currently you should only ask for the log column in the query. Asking for the other column will result in the query not succeeding.<br/>
 The `limit` parameter limits the number of SQL queries that are used to generate the response. Whether you hit the token limit of the provider you use and the cost is something you should be aware of.<br/>
-NOTE: since each milla instance can have its own database, all instances might not necessarily have access to all the data milla is gathering but if you use the same database for all the instances, all instances will have access to all the gathered data.<br/>
+NOTE: since each milla instance can have its own database, all instances might not necessarily have access to all the data milla is gathering. If you use the same database for all the instances, all instances will have access to all the gathered data.<br/>
 
 ### Example Config File
 
@@ -521,10 +531,17 @@ go build
 
 ## Thanks
 
+Milla would not exist without the following projects:
+
 - [girc](https://github.com/lrstanley/girc)
 - [chroma](https://github.com/alecthomas/chroma)
 - [pgx](https://github.com/jackc/pgx)
 - [ollama](https://github.com/ollama/ollama)
+- [toml](https://github.com/BurntSushi/toml)
+
+## TODO
+
+- plugins support
 
 ## Similar Projects
 
diff --git a/main.go b/main.go
index f8f0739..db2a31b 100644
--- a/main.go
+++ b/main.go
@@ -356,19 +356,103 @@ func byteToMByte(bytes uint64,
 	return bytes / 1024 / 1024
 }
 
-func runCommand(
+func handleCustomCommand(
+	args []string,
 	client *girc.Client,
 	event girc.Event,
 	appConfig *TomlConfig,
 ) {
-	cmd := strings.TrimPrefix(event.Last(), appConfig.IrcNick+": ")
-	cmd = strings.TrimSpace(cmd)
-	cmd = strings.TrimPrefix(cmd, "/")
-	args := strings.Split(cmd, " ")
+	log.Println(args)
+	if len(args) < 2 {
+		client.Cmd.Reply(event, errNotEnoughArgs.Error())
+
+		return
+	}
+
+	customCommand := appConfig.CustomCommands[args[1]]
+
+	if customCommand.SQL == "" {
+		client.Cmd.Reply(event, "empty sql commands in the custom command")
+
+		return
+	}
+
+	if appConfig.pool == nil {
+		client.Cmd.Reply(event, "no database connection")
+
+		return
+	}
+
+	log.Println(customCommand.SQL)
+
+	rows, err := appConfig.pool.Query(context.Background(), customCommand.SQL)
+	if err != nil {
+		client.Cmd.Reply(event, "error: "+err.Error())
+
+		return
+	}
+	defer rows.Close()
+
+	logs, err := pgx.CollectRows(rows, pgx.RowToStructByName[LogModel])
+	if err != nil {
+		log.Println(err.Error())
+
+		return
+	}
+
+	log.Println(logs)
+	logs = logs[:customCommand.Limit]
+
+	if err != nil {
+		log.Println(err.Error())
+
+		return
+	}
 
+	switch appConfig.Provider {
+	case "chatgpt":
+		var gptMemory []openai.ChatCompletionMessage
+
+		for _, log := range logs {
+			gptMemory = append(gptMemory, openai.ChatCompletionMessage{
+				Role:    openai.ChatMessageRoleUser,
+				Content: log.Log,
+			})
+		}
+
+		chatGPTRequest(appConfig, client, event, &gptMemory, customCommand.Prompt)
+	case "gemini":
+		var geminiMemory []*genai.Content
+
+		for _, log := range logs {
+			geminiMemory = append(geminiMemory, &genai.Content{
+				Parts: []genai.Part{
+					genai.Text(log.Log),
+				},
+				Role: "user",
+			})
+		}
+
+		geminiRequest(appConfig, client, event, &geminiMemory, customCommand.Prompt)
+	case "ollama":
+		var ollamaMemory []MemoryElement
+
+		for _, log := range logs {
+			ollamaMemory = append(ollamaMemory, MemoryElement{
+				Role:    "user",
+				Content: log.Log,
+			})
+		}
+
+		ollamaRequest(appConfig, client, event, &ollamaMemory, customCommand.Prompt)
+	default:
+	}
+}
+
+func isFromAdmin(admins []string, event girc.Event) bool {
 	messageFromAdmin := false
 
-	for _, admin := range appConfig.Admins {
+	for _, admin := range admins {
 		if event.Source.Name == admin {
 			messageFromAdmin = true
 
@@ -376,7 +460,20 @@ func runCommand(
 		}
 	}
 
-	if !messageFromAdmin {
+	return messageFromAdmin
+}
+
+func runCommand(
+	client *girc.Client,
+	event girc.Event,
+	appConfig *TomlConfig,
+) {
+	cmd := strings.TrimPrefix(event.Last(), appConfig.IrcNick+": ")
+	cmd = strings.TrimSpace(cmd)
+	cmd = strings.TrimPrefix(cmd, "/")
+	args := strings.Split(cmd, " ")
+
+	if appConfig.AdminOnly && !isFromAdmin(appConfig.Admins, event) {
 		return
 	}
 
@@ -431,6 +528,10 @@ func runCommand(
 		client.Cmd.Reply(event, fmt.Sprintf("TotalAlloc: %d MiB", byteToMByte(memStats.TotalAlloc)))
 		client.Cmd.Reply(event, fmt.Sprintf("Sys: %d MiB", byteToMByte(memStats.Sys)))
 	case "join":
+		if !isFromAdmin(appConfig.Admins, event) {
+			break
+		}
+
 		if len(args) < 2 {
 			client.Cmd.Reply(event, errNotEnoughArgs.Error())
 
@@ -439,72 +540,23 @@ func runCommand(
 
 		client.Cmd.Join(args[1])
 	case "leave":
-		if len(args) < 2 {
-			client.Cmd.Reply(event, errNotEnoughArgs.Error())
-
+		if !isFromAdmin(appConfig.Admins, event) {
 			break
 		}
 
-		client.Cmd.Part(args[1])
-	case "cmd":
 		if len(args) < 2 {
 			client.Cmd.Reply(event, errNotEnoughArgs.Error())
 
 			break
 		}
 
-		customCommand := appConfig.CustomCommands[args[1]]
-
-		if customCommand.SQL == "" {
-			client.Cmd.Reply(event, "empty sql commands in the custom command")
-
-			break
-		}
-
-		if appConfig.pool == nil {
-			client.Cmd.Reply(event, "no database connection")
-
-			break
-		}
-
-		log.Println(customCommand.SQL)
-
-		rows, err := appConfig.pool.Query(context.Background(), customCommand.SQL)
-		defer rows.Close()
-
-		if err != nil {
-			client.Cmd.Reply(event, "error: "+err.Error())
-
-			break
-		}
-
-		var gptMemory []openai.ChatCompletionMessage
-
-		logs, err := pgx.CollectRows(rows, pgx.RowToStructByName[LogModel])
-		if err != nil {
-			log.Println(err.Error())
-
-			break
-		}
-
-		log.Println(logs)
-		logs = logs[:customCommand.Limit]
-
-		if err != nil {
-			log.Println(err.Error())
-
+		client.Cmd.Part(args[1])
+	case "cmd":
+		if !isFromAdmin(appConfig.Admins, event) {
 			break
 		}
 
-		for _, log := range logs {
-			gptMemory = append(gptMemory, openai.ChatCompletionMessage{
-				Role:    openai.ChatMessageRoleUser,
-				Content: log.Log,
-			})
-		}
-
-		chatGPTRequest(appConfig, client, event, &gptMemory, customCommand.Prompt)
-
+		handleCustomCommand(args, client, event, appConfig)
 	default:
 		client.Cmd.Reply(event, errUnknCmd.Error())
 	}
@@ -937,22 +989,6 @@ func chatGPTHandler(
 
 func connectToDB(appConfig *TomlConfig, ctx *context.Context, poolChan chan *pgxpool.Pool) {
 	for {
-		if appConfig.DatabaseUser == "" {
-			appConfig.DatabaseUser = os.Getenv("MILLA_DB_USER")
-		}
-
-		if appConfig.DatabasePassword == "" {
-			appConfig.DatabasePassword = os.Getenv("MILLA_DB_PASSWORD")
-		}
-
-		if appConfig.DatabaseAddress == "" {
-			appConfig.DatabaseAddress = os.Getenv("MILLA_DB_ADDRESS")
-		}
-
-		if appConfig.DatabaseName == "" {
-			appConfig.DatabaseName = os.Getenv("MILLA_DB_NAME")
-		}
-
 		dbURL := fmt.Sprintf(
 			"postgres://%s:%s@%s/%s",
 			appConfig.DatabaseUser,
@@ -1059,10 +1095,6 @@ func runIRC(appConfig TomlConfig) {
 		irc.Config.Out = os.Stdout
 	}
 
-	if appConfig.ServerPass == "" {
-		appConfig.ServerPass = os.Getenv("MILLA_SERVER_PASSWORD")
-	}
-
 	irc.Config.ServerPass = appConfig.ServerPass
 
 	if appConfig.Bind != "" {
@@ -1073,17 +1105,7 @@ func runIRC(appConfig TomlConfig) {
 		irc.Config.Name = appConfig.Name
 	}
 
-	saslUser := appConfig.IrcSaslUser
-
-	var saslPass string
-
-	if appConfig.IrcSaslPass == "" {
-		saslPass = os.Getenv("MILLA_SASL_PASSWORD")
-	} else {
-		saslPass = appConfig.IrcSaslPass
-	}
-
-	if appConfig.EnableSasl && saslUser != "" && saslPass != "" {
+	if appConfig.EnableSasl && appConfig.IrcSaslPass != "" && appConfig.IrcSaslUser != "" {
 		irc.Config.SASL = &girc.SASLPlain{
 			User: appConfig.IrcSaslUser,
 			Pass: appConfig.IrcSaslPass,