Implement otel-tg

Signed-off-by: Sid Sun <[email protected]>

cmd/config/config.go

@@ -4,8 +4,11 @@ import (
 	"os"
 )
 
+var GlobalConfig Config
+
 // Config contains all the neccessary configurations
 type Config struct {
+	OpenAIAPI   OpenAI
 	Bot         BotConfig
 	environment string
 }
@@ -17,10 +20,15 @@ func (c Config) GetEnv() string {
 
 // Load reads all config from env to config
 func Load() Config {
-	return Config{
+	GlobalConfig = Config{
 		environment: os.Getenv("APP_ENV"),
 		Bot: BotConfig{
 			tkn: os.Getenv("API_TOKEN"),
 		},
+		OpenAIAPI: OpenAI{
+			Endpoint: os.Getenv("OPENAI_ENDPOINT"),
+			APIKey:   os.Getenv("OPENAI_API_KEY"),
+		},
 	}
+	return GlobalConfig
 }

cmd/config/openai.go

@@ -0,0 +1,6 @@
+package config
+
+type OpenAI struct {
+	Endpoint string
+	APIKey   string
+}

cmd/main.go

@@ -1,12 +1,11 @@
 package main
 
 import (
-	"github.com/sid-sun/sample-bot/cmd/config"
-	"github.com/sid-sun/sample-bot/pkg/bot"
+	"github.com/sid-sun/openwebui-bot/cmd/config"
+	"github.com/sid-sun/openwebui-bot/pkg/bot"
 )
 
 func main() {
 	cfg := config.Load()
-	initLogger(cfg.GetEnv())
-	bot.StartBot(cfg, logger)
+	bot.StartBot(cfg)
 }

go.mod

@@ -1,9 +1,8 @@
-module github.com/sid-sun/sample-bot
+module github.com/sid-sun/openwebui-bot
 
-go 1.13
+go 1.21
 
 require (
+	github.com/bep/debounce v1.2.1
 	github.com/go-telegram-bot-api/telegram-bot-api/v5 v5.5.1
-	go.uber.org/multierr v1.11.0 // indirect
-	go.uber.org/zap v1.26.0
 )

go.sum

@@ -1,66 +1,4 @@
-github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
-github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
-github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/bep/debounce v1.2.1 h1:v67fRdBA9UQu2NhLFXrSg0Brw7CexQekrBwDMM8bzeY=
+github.com/bep/debounce v1.2.1/go.mod h1:H8yggRPQKLUhUoqrJC1bO2xNya7vanpDl7xR3ISbCJ0=
 github.com/go-telegram-bot-api/telegram-bot-api/v5 v5.5.1 h1:wG8n/XJQ07TmjbITcGiUaOtXxdrINDz1b0J1w0SzqDc=
 github.com/go-telegram-bot-api/telegram-bot-api/v5 v5.5.1/go.mod h1:A2S0CWkNylc2phvKXWBBdD3K0iGnDBGbzRpISP2zBl8=
-github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
-github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
-github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
-github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
-github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
-github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
-github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
-github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
-github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
-github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
-github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
-go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk=
-go.uber.org/goleak v1.2.0/go.mod h1:XJYK+MuIchqpmGmUSAzotztawfKvYLUIgg7guXrwVUo=
-go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
-go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
-go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
-go.uber.org/zap v1.26.0 h1:sI7k6L95XOKS281NhVKOFCUNIvv9e0w4BF8N3u+tCRo=
-go.uber.org/zap v1.26.0/go.mod h1:dtElttAiwGvoJ/vj4IwHBS/gXsEu/pZ50mUIRWuG0so=
-golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
-golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs=
-golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
-golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo=
-golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
-golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
-golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
-golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
-golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210510120138-977fb7262007 h1:gG67DSER+11cZvqIMb8S8bt0vZtiN6xWYARwirrOSfE=
-golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
-golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
-golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
-golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
-golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.1.5 h1:ouewzE6p+/VEB31YYnTbEJdi8pFqKp4P4n85vwo3DHA=
-golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
-golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
-gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
-gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/bot/contract/contract.go

@@ -0,0 +1,13 @@
+package contract
+
+type MessageLink struct {
+	Parent   int
+	Children []int
+	Text     string
+	From     string
+}
+
+type CompletionUpdate struct {
+	Message string
+	IsLast  bool
+}

pkg/bot/contract/openai.go

@@ -0,0 +1,49 @@
+package contract
+
+type ChatCompletionPayloadMinimal struct {
+	BasicModelTweaks
+	ModelOptions
+	Messages []ChatMessage `json:"messages"`
+}
+
+type BasicModelTweaks struct {
+	ContextLength int     `json:"context_length"`
+	MaxTokens     int     `json:"max_tokens"`
+	Temperature   float32 `json:"temperature"`
+}
+
+type ChatCompletionPayload struct {
+	ModelOptions
+	ModelTweaks
+	Messages []ChatMessage `json:"messages"`
+}
+
+type ModelOptions struct {
+	Model  string `json:"model"`
+	Stream bool   `json:"stream"`
+}
+
+type ModelTweaks struct {
+	ContextLength    int     `json:"context_length"`
+	MaxTokens        int     `json:"max_tokens"`
+	FrequencyPenalty float32 `json:"frequency_penalty"`
+	PresencePenalty  float32 `json:"presence_penalty"`
+	Temperature      float32 `json:"temperature"`
+	RepeatPenalty    float32 `json:"repeat_penalty"`
+}
+
+type ChatMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type ChatCompletionResponse struct {
+	Choices []ChatCompletionChoice `json:"choices"`
+}
+
+type ChatCompletionChoice struct {
+	Delta struct {
+		Content string `json:"content"`
+	} `json:"delta"`
+	FinishReason string `json:"finish_reason"`
+}

pkg/bot/handlers/completion/completion.go

@@ -0,0 +1,141 @@
+package completion
+
+import (
+	"fmt"
+	"log/slog"
+	"time"
+
+	"github.com/bep/debounce"
+	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
+	"github.com/sid-sun/openwebui-bot/pkg/bot/contract"
+	"github.com/sid-sun/openwebui-bot/pkg/bot/service"
+	"github.com/sid-sun/openwebui-bot/pkg/bot/store"
+)
+
+var logger = slog.Default().With(slog.String("package", "Completion"))
+
+// Handler handles all repeat requests
+func Handler(bot *tgbotapi.BotAPI, update tgbotapi.Update) {
+	logger.Info("[Completion] [Attempt]")
+
+	// Check if this is a resend attempt
+	promptID := update.Message.MessageID
+	isResend := false
+	if update.Message.IsCommand() {
+		switch update.Message.Command() {
+		case "resend":
+			replyToMessage := store.ChatStore[update.FromChat().ID][update.Message.ReplyToMessage.MessageID]
+			if update.Message.ReplyToMessage == nil || replyToMessage == nil {
+				// drop invalid request
+				em := tgbotapi.NewMessage(update.FromChat().ID, "Reply to the message you want to regenerate from - it can't be a previous /resend message")
+				bot.Send(em)
+				return
+			}
+			if replyToMessage.From != update.Message.From.UserName {
+				// invalid request
+				em := tgbotapi.NewMessage(update.FromChat().ID, "Last message for resend must be a user message")
+				bot.Send(em)
+			}
+			promptID = update.Message.ReplyToMessage.MessageID
+			isResend = true
+		}
+	}
+
+	action := tgbotapi.NewChatAction(update.Message.Chat.ID, tgbotapi.ChatTyping)
+	bot.Send(action)
+
+	if !isResend {
+		addMessageToChain(update.Message)
+	}
+
+	updatesChan := make(chan contract.CompletionUpdate, 100)
+	go func() {
+		err := service.GetChatResponseStream(update.Message.Chat.ID, promptID, updatesChan)
+		if err != nil {
+			logger.Error("could not generate completion", slog.String("conext", "GetChatResponseStream"), slog.Any("error", err))
+		}
+	}()
+
+	firstCompletion := <-updatesChan
+
+	msg := tgbotapi.NewMessage(update.Message.Chat.ID, firstCompletion.Message)
+	msg.ReplyToMessageID = promptID
+	// msg.ReplyMarkup = tgbotapi.ForceReply{ForceReply: true}
+
+	botMessage, err := bot.Send(msg)
+	if err != nil {
+		logger.Error("failed to send message", slog.String("context", "new completion message"), slog.Any("error", err))
+		return
+	}
+
+	debounced := debounce.New(5 * time.Millisecond)
+	for completion := range updatesChan {
+		if completion.IsLast {
+			break
+		}
+		edit := tgbotapi.NewEditMessageText(update.Message.Chat.ChatConfig().ChatID, botMessage.MessageID, completion.Message)
+		send := func() {
+			botMessage.Text = edit.Text
+			_, err = bot.Send(edit)
+			if err != nil {
+				logger.Error("failed to send message", slog.String("context", "edit completion message"), slog.Any("error", err))
+				return
+			}
+		}
+		debounced(send)
+	}
+
+	time.Sleep(10 * time.Millisecond)
+	botMessage.Chat = update.Message.Chat
+	botMessage.ReplyToMessage = update.Message
+	if isResend {
+		botMessage.ReplyToMessage = update.Message.ReplyToMessage
+	}
+	addMessageToChain(&botMessage)
+	// To print reply  messages, comment above and uncomment below
+	// bm := addMessageToChain(&botMessage)
+	// printReplyMessages(update.Message.Chat.ID, bm)
+
+	slog.Info("[Completion] [Success]")
+}
+
+func printReplyMessages(chatID int64, m *contract.MessageLink) {
+	if m.Parent != 0 {
+		printReplyMessages(chatID, store.ChatStore[chatID][m.Parent])
+	}
+	fmt.Printf("Message: %s\n", m.Text)
+}
+
+func addMessageToChain(m *tgbotapi.Message) *contract.MessageLink {
+	if store.ChatStore[m.Chat.ID] == nil {
+		store.ChatStore[m.Chat.ID] = make(map[int]*contract.MessageLink)
+	}
+
+	var parent int
+	if m.ReplyToMessage != nil {
+		parent = m.ReplyToMessage.MessageID
+		if store.ChatStore[m.Chat.ID][parent] == nil {
+			addMessageToChain(m.ReplyToMessage)
+		}
+		store.ChatStore[m.Chat.ID][parent].Children = append(store.ChatStore[m.Chat.ID][parent].Children, m.MessageID)
+	}
+
+	// fmt.Printf("Message: %+v\n", m)
+	if m.From == nil {
+		m.From = &tgbotapi.User{UserName: "unknown"}
+	}
+	store.ChatStore[m.Chat.ID][m.MessageID] = &contract.MessageLink{
+		Parent:   parent,
+		Children: []int{},
+		Text:     m.Text,
+		From:     m.From.UserName,
+	}
+
+	// x, err := json.MarshalIndent(store.ChatStore, "", "  ")
+	// if err != nil {
+	// 	fmt.Println(err)
+	// }
+	// fmt.Println(string(x))
+
+	return store.ChatStore[m.Chat.ID][m.MessageID]
+}