From b7185d908c7637163b391e95bd958628c9b59798 Mon Sep 17 00:00:00 2001 From: terminaldweller Date: Tue, 4 Jun 2024 19:14:29 -0400 Subject: added ollama,chatgpt and gemini request functions for the lua plugins to use --- .github/workflows/docker.yaml | 15 ++-- main.go | 169 ++++++++++++++++++++++++------------------ plugins.go | 64 ++++++++++++++-- plugins_test.go | 34 --------- 4 files changed, 164 insertions(+), 118 deletions(-) delete mode 100644 plugins_test.go diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index 33354ce..421f7d2 100644 --- a/.github/workflows/docker.yaml +++ b/.github/workflows/docker.yaml @@ -39,9 +39,14 @@ jobs: tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} provenance: mode=max - # - name: Generate artifact attestation - # uses: actions/attest-build-provenance@v1 + # - name: Docker Scout + # id: docker-scout + # if: ${{ github.event_name == 'pull_request' }} + # uses: docker/scout-action@v1 # with: - # subject-name: index.docker.io/milla - # subject-digest: ${{ steps.push.outputs.digest }} - # push-to-registry: true + # command: cves + # image: ${{ github.event.repository.name }} + # ignore-unchanged: true + # only-severities: critical,high,medium,low + # write-comment: true + # github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/main.go b/main.go index ad5096d..2e4e7be 100644 --- a/main.go +++ b/main.go @@ -328,7 +328,10 @@ func handleCustomCommand( }) } - chatGPTRequest(appConfig, client, event, &gptMemory, customCommand.Prompt) + result := ChatGPTRequestProcessor(appConfig, client, event, &gptMemory, customCommand.Prompt) + if result != "" { + sendToIRC(client, event, result, appConfig.ChromaFormatter) + } case "gemini": var geminiMemory []*genai.Content @@ -341,7 +344,10 @@ func handleCustomCommand( }) } - geminiRequest(appConfig, client, event, &geminiMemory, customCommand.Prompt) + result := GeminiRequestProcessor(appConfig, client, event, &geminiMemory, customCommand.Prompt) + if result != "" { + sendToIRC(client, event, result, appConfig.ChromaFormatter) + } case "ollama": var ollamaMemory []MemoryElement @@ -352,7 +358,10 @@ func handleCustomCommand( }) } - ollamaRequest(appConfig, client, event, &ollamaMemory, customCommand.Prompt) + result := OllamaRequestProcessor(appConfig, client, event, &ollamaMemory, customCommand.Prompt) + if result != "" { + sendToIRC(client, event, result, appConfig.ChromaFormatter) + } default: } } @@ -470,13 +479,12 @@ func runCommand( } } -func doOllamaRequest( +func DoOllamaRequest( appConfig *TomlConfig, client *girc.Client, - event girc.Event, ollamaMemory *[]MemoryElement, prompt string, -) (*http.Response, error) { +) (string, error) { var jsonPayload []byte var err error @@ -504,9 +512,8 @@ func doOllamaRequest( jsonPayload, err = json.Marshal(ollamaRequest) if err != nil { - client.Cmd.ReplyTo(event, "error: "+err.Error()) - return nil, fmt.Errorf("could not marshal json payload: %v", err) + return "", err } log.Printf("json payload: %s", string(jsonPayload)) @@ -516,9 +523,8 @@ func doOllamaRequest( request, err := http.NewRequest(http.MethodPost, appConfig.OllamaEndpoint, bytes.NewBuffer(jsonPayload)) if err != nil { - client.Cmd.ReplyTo(event, "error: "+err.Error()) - return nil, fmt.Errorf("could not make a new http request: %v", err) + return "", err } request = request.WithContext(ctx) @@ -549,66 +555,71 @@ func doOllamaRequest( }, } } + response, err := httpClient.Do(request) - return httpClient.Do(request) -} - -func ollamaRequest( - appConfig *TomlConfig, - client *girc.Client, - event girc.Event, - ollamaMemory *[]MemoryElement, - prompt string, -) { - response, err := doOllamaRequest(appConfig, client, event, ollamaMemory, prompt) - - if response == nil { - return + if err != nil { + return "", err } if err != nil { - client.Cmd.ReplyTo(event, "error: "+err.Error()) - return + return "", err } defer response.Body.Close() log.Println("response body:", response.Body) - var writer bytes.Buffer - var ollamaChatResponse OllamaChatMessagesResponse err = json.NewDecoder(response.Body).Decode(&ollamaChatResponse) + if err != nil { + return "", err + } + + return ollamaChatResponse.Messages.Content, nil +} + +func OllamaRequestProcessor( + appConfig *TomlConfig, + client *girc.Client, + event girc.Event, + ollamaMemory *[]MemoryElement, + prompt string, +) string { + response, err := DoOllamaRequest(appConfig, client, ollamaMemory, prompt) if err != nil { client.Cmd.ReplyTo(event, "error: "+err.Error()) + + return "" } assistantElement := MemoryElement{ Role: "assistant", - Content: ollamaChatResponse.Messages.Content, + Content: response, } *ollamaMemory = append(*ollamaMemory, assistantElement) - log.Println(ollamaChatResponse) + log.Println(response) + + var writer bytes.Buffer err = quick.Highlight(&writer, - ollamaChatResponse.Messages.Content, + response, "markdown", appConfig.ChromaFormatter, appConfig.ChromaStyle) if err != nil { client.Cmd.ReplyTo(event, "error: "+err.Error()) - return + return "" } - sendToIRC(client, event, writer.String(), appConfig.ChromaFormatter) + return writer.String() } -func ollamaHandler( +func OllamaHandler( irc *girc.Client, appConfig *TomlConfig, ollamaMemory *[]MemoryElement, @@ -641,25 +652,26 @@ func ollamaHandler( return } - ollamaRequest(appConfig, client, event, ollamaMemory, prompt) + result := OllamaRequestProcessor(appConfig, client, event, ollamaMemory, prompt) + if result != "" { + sendToIRC(client, event, result, appConfig.ChromaFormatter) + } }) } -func doGeminiRequest( +func DoGeminiRequest( appConfig *TomlConfig, client *girc.Client, - event girc.Event, geminiMemory *[]*genai.Content, prompt string, -) string { +) (string, error) { ctx, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second) defer cancel() clientGemini, err := genai.NewClient(ctx, option.WithAPIKey(appConfig.Apikey)) if err != nil { - client.Cmd.ReplyTo(event, "error: "+err.Error()) - return "" + return "", err } defer clientGemini.Close() @@ -674,22 +686,27 @@ func doGeminiRequest( resp, err := cs.SendMessage(ctx, genai.Text(prompt)) if err != nil { - client.Cmd.ReplyTo(event, "error: "+err.Error()) - return "" + return "", err } - return returnGeminiResponse(resp) + return returnGeminiResponse(resp), nil } -func geminiRequest( +func GeminiRequestProcessor( appConfig *TomlConfig, client *girc.Client, event girc.Event, geminiMemory *[]*genai.Content, prompt string, -) { - geminiResponse := doGeminiRequest(appConfig, client, event, geminiMemory, prompt) +) string { + geminiResponse, err := DoGeminiRequest(appConfig, client, geminiMemory, prompt) + if err != nil { + client.Cmd.ReplyTo(event, "error: "+err.Error()) + + return "" + } + log.Println(geminiResponse) if len(*geminiMemory) > appConfig.MemoryLimit { @@ -712,7 +729,7 @@ func geminiRequest( var writer bytes.Buffer - err := quick.Highlight( + err = quick.Highlight( &writer, geminiResponse, "markdown", @@ -721,13 +738,13 @@ func geminiRequest( if err != nil { client.Cmd.ReplyTo(event, "error: "+err.Error()) - return + return "" } - sendToIRC(client, event, writer.String(), appConfig.ChromaFormatter) + return writer.String() } -func geminiHandler( +func GeminiHandler( irc *girc.Client, appConfig *TomlConfig, geminiMemory *[]*genai.Content, @@ -760,17 +777,20 @@ func geminiHandler( return } - geminiRequest(appConfig, client, event, geminiMemory, prompt) + result := GeminiRequestProcessor(appConfig, client, event, geminiMemory, prompt) + + if result != "" { + sendToIRC(client, event, result, appConfig.ChromaFormatter) + } }) } -func doChatGPTRequest( +func DoChatGPTRequest( appConfig *TomlConfig, client *girc.Client, - event girc.Event, gptMemory *[]openai.ChatCompletionMessage, prompt string, -) (openai.ChatCompletionResponse, error) { +) (string, error) { ctx, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second) defer cancel() @@ -780,17 +800,15 @@ func doChatGPTRequest( proxyURL, err := url.Parse(appConfig.IRCProxy) if err != nil { cancel() - client.Cmd.ReplyTo(event, "error: "+err.Error()) - log.Fatal(err.Error()) + return "", err } dialer, err := proxy.FromURL(proxyURL, &net.Dialer{Timeout: time.Duration(appConfig.RequestTimeout) * time.Second}) if err != nil { cancel() - client.Cmd.ReplyTo(event, "error: "+err.Error()) - log.Fatal(err.Error()) + return "", err } httpClient = http.Client{ @@ -814,27 +832,31 @@ func doChatGPTRequest( Model: appConfig.Model, Messages: *gptMemory, }) + if err != nil { - return resp, err + return "", err + } + + return resp.Choices[0].Message.Content, nil } -func chatGPTRequest( +func ChatGPTRequestProcessor( appConfig *TomlConfig, client *girc.Client, event girc.Event, gptMemory *[]openai.ChatCompletionMessage, prompt string, -) { - resp, err := doChatGPTRequest(appConfig, client, event, gptMemory, prompt) +) string { + resp, err := DoChatGPTRequest(appConfig, client, gptMemory, prompt) if err != nil { client.Cmd.ReplyTo(event, "error: "+err.Error()) - return + return "" } *gptMemory = append(*gptMemory, openai.ChatCompletionMessage{ Role: openai.ChatMessageRoleAssistant, - Content: resp.Choices[0].Message.Content, + Content: resp, }) if len(*gptMemory) > appConfig.MemoryLimit { @@ -845,20 +867,20 @@ func chatGPTRequest( err = quick.Highlight( &writer, - resp.Choices[0].Message.Content, + resp, "markdown", appConfig.ChromaFormatter, appConfig.ChromaStyle) if err != nil { client.Cmd.ReplyTo(event, "error: "+err.Error()) - return + return "" } - sendToIRC(client, event, writer.String(), appConfig.ChromaFormatter) + return writer.String() } -func chatGPTHandler( +func ChatGPTHandler( irc *girc.Client, appConfig *TomlConfig, gptMemory *[]openai.ChatCompletionMessage, @@ -891,7 +913,10 @@ func chatGPTHandler( return } - chatGPTRequest(appConfig, client, event, gptMemory, prompt) + result := ChatGPTRequestProcessor(appConfig, client, event, gptMemory, prompt) + if result != "" { + sendToIRC(client, event, result, appConfig.ChromaFormatter) + } }) } @@ -1039,11 +1064,11 @@ func runIRC(appConfig TomlConfig) { switch appConfig.Provider { case "ollama": - ollamaHandler(irc, &appConfig, &OllamaMemory) + OllamaHandler(irc, &appConfig, &OllamaMemory) case "gemini": - geminiHandler(irc, &appConfig, &GeminiMemory) + GeminiHandler(irc, &appConfig, &GeminiMemory) case "chatgpt": - chatGPTHandler(irc, &appConfig, &GPTMemory) + ChatGPTHandler(irc, &appConfig, &GPTMemory) } go LoadAllPlugins(&appConfig, irc) diff --git a/plugins.go b/plugins.go index 45c2847..0e61c4f 100644 --- a/plugins.go +++ b/plugins.go @@ -7,8 +7,10 @@ import ( "github.com/ailncode/gluaxmlpath" "github.com/cjoudrey/gluahttp" + "github.com/google/generative-ai-go/genai" "github.com/kohkimakimoto/gluayaml" "github.com/lrstanley/girc" + openai "github.com/sashabaranov/go-openai" lua "github.com/yuin/gopher-lua" "gitlab.com/megalithic-llc/gluasocket" ) @@ -210,12 +212,60 @@ func ircPartChannelClosure(luaState *lua.LState, client *girc.Client) func(*lua. } } -func millaModuleLoaderClosure(luaState *lua.LState, client *girc.Client) func(*lua.LState) int { +func ollamaRequestClosure(luaState *lua.LState, client *girc.Client, appConfig *TomlConfig) func(*lua.LState) int { + return func(luaState *lua.LState) int { + prompt := luaState.CheckString(1) + + result, err := DoOllamaRequest(appConfig, client, &[]MemoryElement{}, prompt) + if err != nil { + log.Print(err) + } + + luaState.Push(lua.LString(result)) + + return 1 + } +} + +func geminiRequestClosure(luaState *lua.LState, client *girc.Client, appConfig *TomlConfig) func(*lua.LState) int { + return func(luaState *lua.LState) int { + prompt := luaState.CheckString(1) + + result, err := DoGeminiRequest(appConfig, client, &[]*genai.Content{}, prompt) + if err != nil { + log.Print(err) + } + + luaState.Push(lua.LString(result)) + + return 1 + } +} + +func chatGPTRequestClosure(luaState *lua.LState, client *girc.Client, appConfig *TomlConfig) func(*lua.LState) int { + return func(luaState *lua.LState) int { + prompt := luaState.CheckString(1) + + result, err := DoChatGPTRequest(appConfig, client, &[]openai.ChatCompletionMessage{}, prompt) + if err != nil { + log.Print(err) + } + + luaState.Push(lua.LString(result)) + + return 1 + } +} + +func millaModuleLoaderClosure(luaState *lua.LState, client *girc.Client, appConfig *TomlConfig) func(*lua.LState) int { return func(luaState *lua.LState) int { exports := map[string]lua.LGFunction{ - "send_message": lua.LGFunction(sendMessageClosure(luaState, client)), - "join_channel": lua.LGFunction(ircJoinChannelClosure(luaState, client)), - "part_channel": lua.LGFunction(ircPartChannelClosure(luaState, client)), + "send_message": lua.LGFunction(sendMessageClosure(luaState, client)), + "join_channel": lua.LGFunction(ircJoinChannelClosure(luaState, client)), + "part_channel": lua.LGFunction(ircPartChannelClosure(luaState, client)), + "send_ollama_request": lua.LGFunction(ollamaRequestClosure(luaState, client, appConfig)), + "send_gemini_request": lua.LGFunction(geminiRequestClosure(luaState, client, appConfig)), + "send_chat_gpt_request": lua.LGFunction(chatGPTRequestClosure(luaState, client, appConfig)), } millaModule := luaState.SetFuncs(luaState.NewTable(), exports) @@ -231,11 +281,11 @@ func millaModuleLoaderClosure(luaState *lua.LState, client *girc.Client) func(*l } } -func RunScript(scriptPath string, client *girc.Client) { +func RunScript(scriptPath string, client *girc.Client, appConfig *TomlConfig) { luaState := lua.NewState() defer luaState.Close() - luaState.PreloadModule("milla", millaModuleLoaderClosure(luaState, client)) + luaState.PreloadModule("milla", millaModuleLoaderClosure(luaState, client, appConfig)) gluasocket.Preload(luaState) gluaxmlpath.Preload(luaState) luaState.PreloadModule("http", gluahttp.NewHttpModule(&http.Client{}).Loader) @@ -253,6 +303,6 @@ func LoadAllPlugins(appConfig *TomlConfig, client *girc.Client) { for _, scriptPath := range appConfig.Plugins { log.Print("Loading plugin: ", scriptPath) - go RunScript(scriptPath, client) + go RunScript(scriptPath, client, appConfig) } } diff --git a/plugins_test.go b/plugins_test.go deleted file mode 100644 index 894ff3c..0000000 --- a/plugins_test.go +++ /dev/null @@ -1,34 +0,0 @@ -package main - -import ( - "testing" - - lua "github.com/yuin/gopher-lua" -) - -func TestMetaTable(t *testing.T) { - luaState := lua.NewState() - defer luaState.Close() - - RegisterCustomLuaTypes(luaState) - - if err := luaState.DoString(` - print("Testing MetaTable") - print(toml_config) - - for index, data in ipairs(toml_config) do - print(index, data) - for k,v in pairs(data) do - print("one") - print(k,v) - end - end - - config = toml_config.new() - print(config:IrcServer()) - config:IrcServer("irc.freenode.net") - print(config:IrcServer()) - `); err != nil { - t.Fatal(err) - } -} -- cgit v1.2.3