aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--docker-compose.yaml1
-rw-r--r--main.go64
2 files changed, 49 insertions, 16 deletions
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 19bb2ab..4ffd19c 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -15,6 +15,7 @@ services:
networks:
- millanet
restart: unless-stopped
+ command: ["--config", "/opt/milla/config.toml"]
volumes:
- ./config.toml:/opt/milla/config.toml
cap_drop:
diff --git a/main.go b/main.go
index d7a8fb8..49bb208 100644
--- a/main.go
+++ b/main.go
@@ -5,6 +5,7 @@ import (
"context"
"crypto/tls"
"encoding/json"
+ "flag"
"fmt"
"log"
"net/http"
@@ -56,6 +57,20 @@ type OllamaRequest struct {
Options OllamaRequestOptions `json:"options"`
}
+func printResponse(resp *genai.GenerateContentResponse) string {
+ result := ""
+
+ for _, cand := range resp.Candidates {
+ if cand.Content != nil {
+ for _, part := range cand.Content.Parts {
+ result += fmt.Sprintln(part)
+ }
+ }
+ }
+
+ return result
+}
+
func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
irc := girc.New(girc.Config{
Server: appConfig.IrcServer,
@@ -140,14 +155,17 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
}
var writer bytes.Buffer
- err = quick.Highlight(&writer, ollamaResponse.Response, "markdown", appConfig.ChromaFormatter, appConfig.ChromaStyle)
+ err = quick.Highlight(&writer,
+ ollamaResponse.Response,
+ "markdown",
+ appConfig.ChromaFormatter,
+ appConfig.ChromaStyle)
if err != nil {
client.Cmd.ReplyTo(event, girc.Fmt(fmt.Sprintf("error: %s", err.Error())))
return
}
- fmt.Println(writer.String())
client.Cmd.ReplyTo(event, girc.Fmt("\033[0m"+writer.String()))
}
})
@@ -160,15 +178,22 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second)
defer cancel()
- client_gemini, err := genai.NewClient(ctx, option.WithAPIKey(appConfig.Apikey))
+ transport := http.Transport{
+ Proxy: http.ProxyFromEnvironment,
+ }
+ httpClient := http.Client{
+ Transport: &transport,
+ }
+
+ clientGemini, err := genai.NewClient(ctx, option.WithAPIKey(appConfig.Apikey), option.WithHTTPClient(&httpClient))
if err != nil {
client.Cmd.ReplyTo(event, girc.Fmt(fmt.Sprintf("error: %s", err.Error())))
return
}
- defer client_gemini.Close()
+ defer clientGemini.Close()
- model := client_gemini.GenerativeModel(appConfig.Model)
+ model := clientGemini.GenerativeModel(appConfig.Model)
resp, err := model.GenerateContent(ctx, genai.Text(prompt))
if err != nil {
client.Cmd.ReplyTo(event, girc.Fmt(fmt.Sprintf("error: %s", err.Error())))
@@ -176,18 +201,21 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
return
}
- fmt.Println(resp)
-
- // var writer bytes.Buffer
- // err = quick.Highlight(&writer, resp, "markdown", appConfig.ChromaFormatter, appConfig.ChromaStyle)
- // if err != nil {
- // client.Cmd.ReplyTo(event, girc.Fmt(fmt.Sprintf("error: %s", err.Error())))
+ var writer bytes.Buffer
+ err = quick.Highlight(
+ &writer,
+ printResponse(resp),
+ "markdown",
+ appConfig.ChromaFormatter,
+ appConfig.ChromaStyle)
+ if err != nil {
+ client.Cmd.ReplyTo(event, girc.Fmt(fmt.Sprintf("error: %s", err.Error())))
- // return
- // }
+ return
+ }
- // fmt.Println(writer.String())
- // client.Cmd.ReplyTo(event, girc.Fmt("\033[0m"+writer.String()))
+ fmt.Println(writer.String())
+ client.Cmd.ReplyTo(event, girc.Fmt("\033[0m"+writer.String()))
}
})
}
@@ -208,7 +236,11 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
func main() {
var appConfig TomlConfig
- data, err := os.ReadFile("/opt/milla/config.toml")
+ configPath := flag.String("config", "./config-gemini.toml", "path to the config file")
+
+ flag.Parse()
+
+ data, err := os.ReadFile(*configPath)
if err != nil {
log.Fatal(err)
}