fixes #33

4 files changed, 55 insertions, 18 deletions

diff --git a/README.md b/README.md
index cbc3ff0..393ea5b 100644
--- a/README.md
+++ b/README.md
@@ -2,13 +2,7 @@
 
 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 providers:
-
-* Ollama
-* Openai
-* Gemini
-
+- sends things over to an LLM when you ask it questions and prints the answer with optional syntax-highlighting.Currently supported providers: Ollama, Openai, Gemini <br/>
 - 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/>
@@ -247,13 +241,20 @@ Custom commands let you define a command that does a SQL query to the database a
 
 ```toml
 [ircd.devinet_terra.customCommands.digest]
-sql = "select log from liberanet_milla_us_market_news;"
-limit = 10
-prompt = "give me digest of the provided news"
+sql = "select log from liberanet_milla_us_market_news order by log desc;"
+limit = 300
+context = ["you are a sentiment-analysis bot"]
+prompt= "i have provided to you news headlines in the form of previous conversations between you and me using the user role. please provide the digest of the news for me."
 [ircd.devinet_terra.customCommands.summarize]
-sql= "select log from liberanet_milla_us_market_news;"
+sql= "select log from liberanet_milla_us_market_news order by log desc;"
 limit= 300
-prompt= "given all the data, summarize the news for me"
+context = ["you are a sentiment-analysis bot"]
+prompt= "i have provided to you news headlines in the form of previous conversations between you and me using the user role. please summarize the provided news for me. provide some details."
+[ircd.devinet_terra.customCommands.canada]
+sql= "select log from liberanet_milla_us_market_news order by log desc;"
+limit= 300
+context = ["you are a canadian news anchor", "you only care about news that is relevant to canada"]
+prompt= "i have provided to you news headlines in the form of previous conversations between you and me using the user role. please summarize the provided news for me. provide some details."
 ```
 
 In the above example digest and summarize will be the names of the commands: `milla: /cmd summarize`.<br/>
diff --git a/main.go b/main.go
index b7f47e3..f8dd3bb 100644
--- a/main.go
+++ b/main.go
@@ -113,6 +113,15 @@ func getTableFromChanName(channel, ircdName string) string {
 	return tableName
 }
 
+func stripColorCodes(input string) string {
+	re := regexp.MustCompile(`\x1b\[[0-9;]*m`)
+	input = re.ReplaceAllString(input, "")
+	re = regexp.MustCompile(`\x03(?:\d{1,2}(?:,\d{1,2})?)?`)
+	input = re.ReplaceAllString(input, "")
+
+	return input
+}
+
 func sanitizeLog(log string) string {
 	sanitizeLog := strings.ReplaceAll(log, "'", " ")
 
@@ -328,6 +337,18 @@ func handleCustomCommand(
 			})
 		}
 
+		for _, customContext := range customCommand.Context {
+			gptMemory = append(gptMemory, openai.ChatCompletionMessage{
+				Role:    openai.ChatMessageRoleUser,
+				Content: customContext,
+			})
+		}
+
+		var bigPrompt string
+		for _, log := range logs {
+			bigPrompt += log.Log + "\n"
+		}
+
 		result := ChatGPTRequestProcessor(appConfig, client, event, &gptMemory, customCommand.Prompt)
 		if result != "" {
 			sendToIRC(client, event, result, appConfig.ChromaFormatter)
@@ -344,6 +365,15 @@ func handleCustomCommand(
 			})
 		}
 
+		for _, customContext := range customCommand.Context {
+			geminiMemory = append(geminiMemory, &genai.Content{
+				Parts: []genai.Part{
+					genai.Text(customContext),
+				},
+				Role: "user",
+			})
+		}
+
 		result := GeminiRequestProcessor(appConfig, client, event, &geminiMemory, customCommand.Prompt)
 		if result != "" {
 			sendToIRC(client, event, result, appConfig.ChromaFormatter)
@@ -358,6 +388,13 @@ func handleCustomCommand(
 			})
 		}
 
+		for _, customContext := range customCommand.Context {
+			ollamaMemory = append(ollamaMemory, MemoryElement{
+				Role:    "user",
+				Content: customContext,
+			})
+		}
+
 		result := OllamaRequestProcessor(appConfig, client, event, &ollamaMemory, customCommand.Prompt)
 		if result != "" {
 			sendToIRC(client, event, result, appConfig.ChromaFormatter)
@@ -996,7 +1033,7 @@ func scrapeChannel(irc *girc.Client, poolChan chan *pgxpool.Pool, appConfig Toml
 			"insert into %s (channel,log,nick) values ('%s','%s','%s')",
 			tableName,
 			sanitizeLog(event.Params[0]),
-			event.Last(),
+			stripColorCodes(event.Last()),
 			event.Source.Name,
 		)
 
diff --git a/plugins/rss.lua b/plugins/rss.lua
index aa9c757..880a0fd 100644
--- a/plugins/rss.lua
+++ b/plugins/rss.lua
@@ -1,5 +1,3 @@
-#!/usr/bin/env lua5.1
-
 local milla = require("milla")
 local yaml = require("yaml")
 local http = require("http")
diff --git a/types.go b/types.go
index 3b7f3ff..98738bd 100644
--- a/types.go
+++ b/types.go
@@ -17,9 +17,10 @@ type LogModel struct {
 }
 
 type CustomCommand struct {
-	SQL    string `toml:"sql"`
-	Limit  int    `toml:"limit"`
-	Prompt string `toml:"prompt"`
+	SQL     string   `toml:"sql"`
+	Limit   int      `toml:"limit"`
+	Context []string `toml:"context"`
+	Prompt  string   `toml:"prompt"`
 }
 
 type LuaLstates struct {