From 4736e43631ed21fd14741daa1dde746687d330fa Mon Sep 17 00:00:00 2001
From: Grail Finder <wohilas@gmail.com>
Date: Sat, 4 Jan 2025 18:13:13 +0300
Subject: Feat (RAG): tying tui calls to rag funcs [WIP; skip-ci]

RAG itself is annoying to properly implement, plucking sentences with no
context is useless. Also it should not be a part of main package, same
for goes for tui. The number of global vars is absurd.
---
 bot.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

(limited to 'bot.go')

diff --git a/bot.go b/bot.go
index 3ee6277..41ae43b 100644
--- a/bot.go
+++ b/bot.go
@@ -15,6 +15,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/neurosnap/sentences/english"
 	"github.com/rivo/tview"
 )
 
@@ -40,6 +41,16 @@ func formMsg(chatBody *models.ChatBody, newMsg, role string) io.Reader {
 	if newMsg != "" { // otherwise let the bot continue
 		newMsg := models.RoleMsg{Role: role, Content: newMsg}
 		chatBody.Messages = append(chatBody.Messages, newMsg)
+		// if rag
+		if cfg.RAGEnabled {
+			ragResp, err := chatRagUse(newMsg.Content)
+			if err != nil {
+				logger.Error("failed to form a rag msg", "error", err)
+				return nil
+			}
+			ragMsg := models.RoleMsg{Role: cfg.ToolRole, Content: ragResp}
+			chatBody.Messages = append(chatBody.Messages, ragMsg)
+		}
 	}
 	data, err := json.Marshal(chatBody)
 	if err != nil {
@@ -107,6 +118,40 @@ func sendMsgToLLM(body io.Reader) {
 	}
 }
 
+func chatRagUse(qText string) (string, error) {
+	tokenizer, err := english.NewSentenceTokenizer(nil)
+	if err != nil {
+		return "", err
+	}
+	// TODO: this where llm should find the questions in text and ask them
+	questionsS := tokenizer.Tokenize(qText)
+	questions := make([]string, len(questionsS))
+	for i, q := range questionsS {
+		questions[i] = q.Text
+	}
+	respVecs := []*models.VectorRow{}
+	for i, q := range questions {
+		emb, err := lineToVector(q)
+		if err != nil {
+			logger.Error("failed to get embs", "error", err, "index", i, "question", q)
+			continue
+		}
+		vec, err := searchEmb(emb)
+		if err != nil {
+			logger.Error("failed to get embs", "error", err, "index", i, "question", q)
+			continue
+		}
+		respVecs = append(respVecs, vec)
+		// logger.Info("returned vector from query search", "question", q, "vec", vec)
+	}
+	// get raw text
+	resps := []string{}
+	for _, rv := range respVecs {
+		resps = append(resps, rv.RawText)
+	}
+	return strings.Join(resps, "\n"), nil
+}
+
 func chatRound(userMsg, role string, tv *tview.TextView, regen bool) {
 	botRespMode = true
 	reader := formMsg(chatBody, userMsg, role)
@@ -294,4 +339,5 @@ func init() {
 		Stream:   true,
 		Messages: lastChat,
 	}
+	// tempLoad()
 }
-- 
cgit v1.2.3