1 files changed, 208 insertions, 0 deletions
diff --git a/ollama.go b/ollama.go
new file mode 100644
index 0000000..d583499
--- /dev/null
+++ b/ollama.go
@@ -0,0 +1,208 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"log"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+
+	"github.com/alecthomas/chroma/v2/quick"
+	"github.com/lrstanley/girc"
+	"golang.org/x/net/proxy"
+)
+
+func DoOllamaRequest(
+	appConfig *TomlConfig,
+	ollamaMemory *[]MemoryElement,
+	prompt, systemPrompt string,
+) (string, error) {
+	var jsonPayload []byte
+
+	var err error
+
+	memoryElement := MemoryElement{
+		Role:    "user",
+		Content: prompt,
+	}
+
+	if len(*ollamaMemory) > appConfig.MemoryLimit {
+		*ollamaMemory = []MemoryElement{}
+
+		for _, context := range appConfig.Context {
+			*ollamaMemory = append(*ollamaMemory, MemoryElement{
+				Role:    "assistant",
+				Content: context,
+			})
+		}
+	}
+
+	*ollamaMemory = append(*ollamaMemory, memoryElement)
+
+	ollamaRequest := OllamaChatRequest{
+		Model:     appConfig.Model,
+		KeepAlive: time.Duration(appConfig.KeepAlive),
+		Stream:    false,
+		Messages:  *ollamaMemory,
+		System:    systemPrompt,
+		Options: OllamaRequestOptions{
+			Mirostat:      appConfig.OllamaMirostat,
+			MirostatEta:   appConfig.OllamaMirostatEta,
+			MirostatTau:   appConfig.OllamaMirostatTau,
+			NumCtx:        appConfig.OllamaNumCtx,
+			RepeatLastN:   appConfig.OllamaRepeatLastN,
+			RepeatPenalty: appConfig.OllamaRepeatPenalty,
+			Temperature:   appConfig.Temperature,
+			Seed:          appConfig.OllamaSeed,
+			NumPredict:    appConfig.OllamaNumPredict,
+			TopK:          appConfig.TopK,
+			TopP:          appConfig.TopP,
+			MinP:          appConfig.OllamaMinP,
+		},
+	}
+
+	jsonPayload, err = json.Marshal(ollamaRequest)
+	if err != nil {
+		return "", err
+	}
+
+	log.Printf("json payload: %s", string(jsonPayload))
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second)
+	defer cancel()
+
+	request, err := http.NewRequest(http.MethodPost, appConfig.Endpoint, bytes.NewBuffer(jsonPayload))
+	if err != nil {
+		return "", err
+	}
+
+	request = request.WithContext(ctx)
+	request.Header.Set("Content-Type", "application/json")
+
+	var httpClient http.Client
+
+	var dialer proxy.Dialer
+
+	if appConfig.LLMProxy != "" {
+		proxyURL, err := url.Parse(appConfig.IRCProxy)
+		if err != nil {
+			cancel()
+
+			log.Fatal(err.Error())
+		}
+
+		dialer, err = proxy.FromURL(proxyURL, &net.Dialer{Timeout: time.Duration(appConfig.RequestTimeout) * time.Second})
+		if err != nil {
+			cancel()
+
+			log.Fatal(err.Error())
+		}
+
+		httpClient = http.Client{
+			Transport: &http.Transport{
+				Dial: dialer.Dial,
+			},
+		}
+	}
+
+	response, err := httpClient.Do(request)
+	if err != nil {
+		return "", err
+	}
+
+	defer response.Body.Close()
+
+	var ollamaChatResponse OllamaChatMessagesResponse
+
+	err = json.NewDecoder(response.Body).Decode(&ollamaChatResponse)
+	if err != nil {
+		return "", err
+	}
+
+	log.Println("ollama chat response: ", ollamaChatResponse)
+
+	return ollamaChatResponse.Messages.Content, nil
+}
+
+func OllamaRequestProcessor(
+	appConfig *TomlConfig,
+	client *girc.Client,
+	event girc.Event,
+	ollamaMemory *[]MemoryElement,
+	prompt, systemPrompt string,
+) string {
+	response, err := DoOllamaRequest(appConfig, ollamaMemory, prompt, systemPrompt)
+	if err != nil {
+		client.Cmd.ReplyTo(event, "error: "+err.Error())
+
+		return ""
+	}
+
+	assistantElement := MemoryElement{
+		Role:    "assistant",
+		Content: response,
+	}
+
+	*ollamaMemory = append(*ollamaMemory, assistantElement)
+
+	log.Println(response)
+
+	var writer bytes.Buffer
+
+	err = quick.Highlight(&writer,
+		response,
+		"markdown",
+		appConfig.ChromaFormatter,
+		appConfig.ChromaStyle)
+	if err != nil {
+		client.Cmd.ReplyTo(event, "error: "+err.Error())
+
+		return ""
+	}
+
+	return writer.String()
+}
+
+func OllamaHandler(
+	irc *girc.Client,
+	appConfig *TomlConfig,
+	ollamaMemory *[]MemoryElement,
+) {
+	irc.Handlers.AddBg(girc.PRIVMSG, func(client *girc.Client, event girc.Event) {
+		if !strings.HasPrefix(event.Last(), appConfig.IrcNick+": ") {
+			return
+		}
+
+		if appConfig.AdminOnly {
+			byAdmin := false
+
+			for _, admin := range appConfig.Admins {
+				if event.Source.Name == admin {
+					byAdmin = true
+				}
+			}
+
+			if !byAdmin {
+				return
+			}
+		}
+
+		prompt := strings.TrimPrefix(event.Last(), appConfig.IrcNick+": ")
+		log.Println(prompt)
+
+		if string(prompt[0]) == "/" {
+			runCommand(client, event, appConfig)
+
+			return
+		}
+
+		result := OllamaRequestProcessor(appConfig, client, event, ollamaMemory, prompt, appConfig.SystemPrompt)
+		if result != "" {
+			SendToIRC(client, event, result, appConfig.ChromaFormatter)
+		}
+	})
+}