diff options
| -rw-r--r-- | docker-compose.yaml | 1 | ||||
| -rw-r--r-- | main.go | 64 |
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: @@ -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) } |