Refactor: rag to sep package

bot.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"elefant/config"
 	"elefant/models"
+	"elefant/rag"
 	"elefant/storage"
 	"encoding/json"
 	"fmt"
@@ -33,6 +34,7 @@ var (
 	defaultStarter      = []models.RoleMsg{}
 	defaultStarterBytes = []byte{}
 	interruptResp       = false
+	ragger              *rag.RAG
 )
 
 // ====
@@ -129,26 +131,34 @@ func chatRagUse(qText string) (string, error) {
 	for i, q := range questionsS {
 		questions[i] = q.Text
 	}
-	respVecs := []*models.VectorRow{}
+	respVecs := []models.VectorRow{}
 	for i, q := range questions {
-		emb, err := lineToVector(q)
+		emb, err := ragger.LineToVector(q)
 		if err != nil {
 			logger.Error("failed to get embs", "error", err, "index", i, "question", q)
 			continue
 		}
-		vec, err := searchEmb(emb)
+		// e := &models.EmbeddingResp{
+		// 	Embedding: emb,
+		// }
+		// vecs, err := ragger.SearchEmb(e)
+		vecs, err := store.SearchClosest(emb)
 		if err != nil {
-			logger.Error("failed to get embs", "error", err, "index", i, "question", q)
+			logger.Error("failed to query embs", "error", err, "index", i, "question", q)
 			continue
 		}
-		respVecs = append(respVecs, vec)
+		respVecs = append(respVecs, vecs...)
 		// logger.Info("returned vector from query search", "question", q, "vec", vec)
 	}
 	// get raw text
 	resps := []string{}
+	logger.Info("sqlvec resp", "vecs", respVecs)
 	for _, rv := range respVecs {
 		resps = append(resps, rv.RawText)
 	}
+	if len(resps) == 0 {
+		return "No related results from vector storage.", nil
+	}
 	return strings.Join(resps, "\n"), nil
 }
 
@@ -326,6 +336,7 @@ func init() {
 	if store == nil {
 		os.Exit(1)
 	}
+	ragger = rag.New(logger, store, cfg)
 	// https://github.com/coreydaley/ggerganov-llama.cpp/blob/master/examples/server/README.md
 	// load all chats in memory
 	if _, err := loadHistoryChats(); err != nil {

main.go

@@ -10,7 +10,7 @@ var (
 	botRespMode   = false
 	editMode      = false
 	selectedIndex = int(-1)
-	indexLine     = "F12 to show keys help; bot resp mode: %v; char: %s; chat: %s; RAGEnabled: %v"
+	indexLine     = "F12 to show keys help; bot resp mode: %v; char: %s; chat: %s; RAGEnabled: %v; EmbedURL: %s"
 	focusSwitcher = map[tview.Primitive]tview.Primitive{}
 )
 

rag.go

@@ -1,222 +0,0 @@
-package main
-
-import (
-	"bytes"
-	"context"
-	"elefant/models"
-	"encoding/json"
-	"errors"
-	"fmt"
-	"net/http"
-	"os"
-
-	"github.com/neurosnap/sentences/english"
-)
-
-func loadRAG(fpath string) error {
-	data, err := os.ReadFile(fpath)
-	if err != nil {
-		return err
-	}
-	fileText := string(data)
-	tokenizer, err := english.NewSentenceTokenizer(nil)
-	if err != nil {
-		return err
-	}
-	sentences := tokenizer.Tokenize(fileText)
-	sents := make([]string, len(sentences))
-	for i, s := range sentences {
-		sents[i] = s.Text
-	}
-	var (
-		// TODO: to config
-		workers   = 5
-		batchSize = 200
-		//
-		left     = 0
-		right    = batchSize
-		batchCh  = make(chan map[int][]string)
-		vectorCh = make(chan []models.VectorRow)
-		errCh    = make(chan error)
-	)
-	if len(sents) < batchSize {
-		batchSize = len(sents)
-	}
-	// fill input channel
-	for {
-		if right > len(sents) {
-			batchCh <- map[int][]string{left: sents[left:]}
-			break
-		}
-		batchCh <- map[int][]string{left: sents[left:right]}
-		left, right = right, right+batchSize
-	}
-	// TODO: cancel complains, replace ctx with done chan
-	ctx, cancel := context.WithCancel(context.Background())
-	for w := 0; w < workers; w++ {
-		go batchToVectorHFAsync(ctx, cancel, len(sents), batchCh, vectorCh, errCh)
-	}
-	// write to db
-	return writeVectors(vectorCh)
-}
-
-func writeVectors(vectorCh <-chan []models.VectorRow) error {
-	for batch := range vectorCh {
-		for _, vector := range batch {
-			if err := store.WriteVector(&vector); err != nil {
-				return err
-			}
-		}
-	}
-	return nil
-}
-
-func batchToVectorHFAsync(ctx context.Context, close context.CancelFunc, limit int,
-	inputCh <-chan map[int][]string, vectorCh chan<- []models.VectorRow, errCh chan error) {
-	for {
-		select {
-		case linesMap := <-inputCh:
-			for leftI, v := range linesMap {
-				FecthEmbHF(v, errCh, vectorCh, fmt.Sprintf("test_%d", leftI))
-				if leftI+200 >= limit { // last batch
-					close()
-					return
-				}
-			}
-		case <-ctx.Done():
-			logger.Error("got ctx done")
-			return
-		case err := <-errCh:
-			logger.Error("got an error", "error", err)
-			close()
-			return
-		}
-	}
-}
-
-func FecthEmbHF(lines []string, errCh chan error, vectorCh chan<- []models.VectorRow, slug string) {
-	payload, err := json.Marshal(
-		map[string]any{"inputs": lines, "options": map[string]bool{"wait_for_model": true}},
-	)
-	if err != nil {
-		logger.Error("failed to marshal payload", "err:", err.Error())
-		errCh <- err
-		return
-	}
-	req, err := http.NewRequest("POST", cfg.EmbedURL, bytes.NewReader(payload))
-	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", cfg.HFToken))
-	resp, err := httpClient.Do(req)
-	// nolint
-	// resp, err := httpClient.Post(cfg.EmbedURL, "application/json", bytes.NewReader(payload))
-	if err != nil {
-		logger.Error("failed to embedd line", "err:", err.Error())
-		errCh <- err
-		return
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != 200 {
-		logger.Error("non 200 resp", "code", resp.StatusCode)
-		errCh <- err
-		return
-	}
-	emb := [][]float32{}
-	if err := json.NewDecoder(resp.Body).Decode(&emb); err != nil {
-		logger.Error("failed to embedd line", "err:", err.Error())
-		errCh <- err
-		return
-	}
-	if len(emb) == 0 {
-		logger.Error("empty emb")
-		err = errors.New("empty emb")
-		errCh <- err
-		return
-	}
-	vectors := make([]models.VectorRow, len(emb))
-	for i, e := range emb {
-		vector := models.VectorRow{
-			Embeddings: e,
-			RawText:    lines[i],
-			Slug:       slug,
-		}
-		vectors[i] = vector
-	}
-	vectorCh <- vectors
-}
-
-func batchToVectorHF(lines []string) ([][]float32, error) {
-	payload, err := json.Marshal(
-		map[string]any{"inputs": lines, "options": map[string]bool{"wait_for_model": true}},
-	)
-	if err != nil {
-		logger.Error("failed to marshal payload", "err:", err.Error())
-		return nil, err
-	}
-	req, err := http.NewRequest("POST", cfg.EmbedURL, bytes.NewReader(payload))
-	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", cfg.HFToken))
-	resp, err := httpClient.Do(req)
-	// nolint
-	// resp, err := httpClient.Post(cfg.EmbedURL, "application/json", bytes.NewReader(payload))
-	if err != nil {
-		logger.Error("failed to embedd line", "err:", err.Error())
-		return nil, err
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != 200 {
-		logger.Error("non 200 resp", "code", resp.StatusCode)
-		return nil, err
-	}
-	emb := [][]float32{}
-	if err := json.NewDecoder(resp.Body).Decode(&emb); err != nil {
-		logger.Error("failed to embedd line", "err:", err.Error())
-		return nil, err
-	}
-	if len(emb) == 0 {
-		logger.Error("empty emb")
-		err = errors.New("empty emb")
-		return nil, err
-	}
-	return emb, nil
-}
-
-func lineToVector(line string) (*models.EmbeddingResp, error) {
-	payload, err := json.Marshal(map[string]string{"content": line})
-	if err != nil {
-		logger.Error("failed to marshal payload", "err:", err.Error())
-		return nil, err
-	}
-	// nolint
-	resp, err := httpClient.Post(cfg.EmbedURL, "application/json", bytes.NewReader(payload))
-	if err != nil {
-		logger.Error("failed to embedd line", "err:", err.Error())
-		return nil, err
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != 200 {
-		logger.Error("non 200 resp", "code", resp.StatusCode)
-		return nil, err
-	}
-	emb := models.EmbeddingResp{}
-	if err := json.NewDecoder(resp.Body).Decode(&emb); err != nil {
-		logger.Error("failed to embedd line", "err:", err.Error())
-		return nil, err
-	}
-	if len(emb.Embedding) == 0 {
-		logger.Error("empty emb")
-		err = errors.New("empty emb")
-		return nil, err
-	}
-	return &emb, nil
-}
-
-func saveLine(topic, line string, emb *models.EmbeddingResp) error {
-	row := &models.VectorRow{
-		Embeddings: emb.Embedding,
-		Slug:       topic,
-		RawText:    line,
-	}
-	return store.WriteVector(row)
-}
-
-func searchEmb(emb *models.EmbeddingResp) (*models.VectorRow, error) {
-	return store.SearchClosest(emb.Embedding)
-}

rag/main.go

@@ -0,0 +1,240 @@
+package rag
+
+import (
+	"bytes"
+	"elefant/config"
+	"elefant/models"
+	"elefant/storage"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"os"
+
+	"github.com/neurosnap/sentences/english"
+)
+
+type RAG struct {
+	logger *slog.Logger
+	store  storage.FullRepo
+	cfg    *config.Config
+}
+
+func New(l *slog.Logger, s storage.FullRepo, cfg *config.Config) *RAG {
+	return &RAG{
+		logger: l,
+		store:  s,
+		cfg:    cfg,
+	}
+}
+
+func (r *RAG) LoadRAG(fpath string) error {
+	data, err := os.ReadFile(fpath)
+	if err != nil {
+		return err
+	}
+	r.logger.Info("rag: loaded file", "fp", fpath)
+	fileText := string(data)
+	tokenizer, err := english.NewSentenceTokenizer(nil)
+	if err != nil {
+		return err
+	}
+	sentences := tokenizer.Tokenize(fileText)
+	sents := make([]string, len(sentences))
+	r.logger.Info("rag: sentences", "#", len(sents))
+	for i, s := range sentences {
+		sents[i] = s.Text
+	}
+	// TODO: maybe better to decide batch size based on sentences len
+	var (
+		// TODO: to config
+		workers   = 5
+		batchSize = 200
+		maxChSize = 1000
+		//
+		left     = 0
+		right    = batchSize
+		batchCh  = make(chan map[int][]string, maxChSize)
+		vectorCh = make(chan []models.VectorRow, maxChSize)
+		errCh    = make(chan error, 1)
+		doneCh   = make(chan bool, 1)
+	)
+	if len(sents) < batchSize {
+		batchSize = len(sents)
+	}
+	// fill input channel
+	ctn := 0
+	for {
+		if right > len(sents) {
+			batchCh <- map[int][]string{left: sents[left:]}
+			break
+		}
+		batchCh <- map[int][]string{left: sents[left:right]}
+		left, right = right, right+batchSize
+		ctn++
+	}
+	r.logger.Info("finished batching", "batches#", len(batchCh))
+	for w := 0; w < workers; w++ {
+		go r.batchToVectorHFAsync(len(sents), batchCh, vectorCh, errCh, doneCh)
+	}
+	// write to db
+	return r.writeVectors(vectorCh, doneCh)
+}
+
+func (r *RAG) writeVectors(vectorCh <-chan []models.VectorRow, doneCh <-chan bool) error {
+	for {
+		select {
+		case batch := <-vectorCh:
+			for _, vector := range batch {
+				if err := r.store.WriteVector(&vector); err != nil {
+					return err
+				}
+			}
+			r.logger.Info("wrote batch to db", "size", len(batch))
+		case <-doneCh:
+			r.logger.Info("rag finished")
+			return nil
+		}
+	}
+}
+
+func (r *RAG) batchToVectorHFAsync(limit int, inputCh <-chan map[int][]string,
+	vectorCh chan<- []models.VectorRow, errCh chan error, doneCh chan bool) {
+	r.logger.Info("to vector batches", "batches#", len(inputCh))
+	for {
+		select {
+		case linesMap := <-inputCh:
+			// r.logger.Info("batch from ch")
+			for leftI, v := range linesMap {
+				// r.logger.Info("fetching", "index", leftI)
+				r.fecthEmbHF(v, errCh, vectorCh, fmt.Sprintf("test_%d", leftI))
+				if leftI+200 >= limit { // last batch
+					doneCh <- true
+					return
+				}
+				// r.logger.Info("done feitching", "index", leftI)
+			}
+		case <-doneCh:
+			r.logger.Info("got done")
+			close(errCh)
+			close(doneCh)
+			return
+		case err := <-errCh:
+			r.logger.Error("got an error", "error", err)
+			return
+		}
+	}
+}
+
+func (r *RAG) fecthEmbHF(lines []string, errCh chan error, vectorCh chan<- []models.VectorRow, slug string) {
+	payload, err := json.Marshal(
+		map[string]any{"inputs": lines, "options": map[string]bool{"wait_for_model": true}},
+	)
+	if err != nil {
+		r.logger.Error("failed to marshal payload", "err:", err.Error())
+		errCh <- err
+		return
+	}
+	req, err := http.NewRequest("POST", r.cfg.EmbedURL, bytes.NewReader(payload))
+	if err != nil {
+		r.logger.Error("failed to create new req", "err:", err.Error())
+		errCh <- err
+		return
+	}
+	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", r.cfg.HFToken))
+	resp, err := http.DefaultClient.Do(req)
+	// nolint
+	// resp, err := httpClient.Post(cfg.EmbedURL, "application/json", bytes.NewReader(payload))
+	if err != nil {
+		r.logger.Error("failed to embedd line", "err:", err.Error())
+		errCh <- err
+		return
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != 200 {
+		r.logger.Error("non 200 resp", "code", resp.StatusCode)
+		return
+		// err = fmt.Errorf("non 200 resp; url: %s; code %d", r.cfg.EmbedURL, resp.StatusCode)
+		// errCh <- err
+		// return
+	}
+	emb := [][]float32{}
+	if err := json.NewDecoder(resp.Body).Decode(&emb); err != nil {
+		r.logger.Error("failed to embedd line", "err:", err.Error())
+		errCh <- err
+		return
+	}
+	if len(emb) == 0 {
+		r.logger.Error("empty emb")
+		err = errors.New("empty emb")
+		errCh <- err
+		return
+	}
+	vectors := make([]models.VectorRow, len(emb))
+	for i, e := range emb {
+		vector := models.VectorRow{
+			Embeddings: e,
+			RawText:    lines[i],
+			Slug:       slug,
+		}
+		vectors[i] = vector
+	}
+	vectorCh <- vectors
+}
+
+func (r *RAG) LineToVector(line string) ([]float32, error) {
+	// payload, err := json.Marshal(map[string]string{"content": line})
+	lines := []string{line}
+	payload, err := json.Marshal(
+		map[string]any{"inputs": lines, "options": map[string]bool{"wait_for_model": true}},
+	)
+	if err != nil {
+		r.logger.Error("failed to marshal payload", "err:", err.Error())
+		return nil, err
+	}
+	// nolint
+	req, err := http.NewRequest("POST", r.cfg.EmbedURL, bytes.NewReader(payload))
+	if err != nil {
+		r.logger.Error("failed to create new req", "err:", err.Error())
+		return nil, err
+	}
+	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", r.cfg.HFToken))
+	resp, err := http.DefaultClient.Do(req)
+	// resp, err := req.Post(r.cfg.EmbedURL, "application/json", bytes.NewReader(payload))
+	if err != nil {
+		r.logger.Error("failed to embedd line", "err:", err.Error())
+		return nil, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != 200 {
+		err = fmt.Errorf("non 200 resp; code: %v\n", resp.StatusCode)
+		r.logger.Error(err.Error())
+		return nil, err
+	}
+	// emb := models.EmbeddingResp{}
+	emb := [][]float32{}
+	if err := json.NewDecoder(resp.Body).Decode(&emb); err != nil {
+		r.logger.Error("failed to embedd line", "err:", err.Error())
+		return nil, err
+	}
+	if len(emb) == 0 || len(emb[0]) == 0 {
+		r.logger.Error("empty emb")
+		err = errors.New("empty emb")
+		return nil, err
+	}
+	return emb[0], nil
+}
+
+func (r *RAG) saveLine(topic, line string, emb *models.EmbeddingResp) error {
+	row := &models.VectorRow{
+		Embeddings: emb.Embedding,
+		Slug:       topic,
+		RawText:    line,
+	}
+	return r.store.WriteVector(row)
+}
+
+func (r *RAG) SearchEmb(emb *models.EmbeddingResp) ([]models.VectorRow, error) {
+	return r.store.SearchClosest(emb.Embedding)
+}

storage/vector.go

@@ -4,7 +4,6 @@ import (
 	"elefant/models"
 	"errors"
 	"fmt"
-	"log"
 	"unsafe"
 
 	sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/ncruces"
@@ -12,7 +11,7 @@ import (
 
 type VectorRepo interface {
 	WriteVector(*models.VectorRow) error
-	SearchClosest(q []float32) (*models.VectorRow, error)
+	SearchClosest(q []float32) ([]models.VectorRow, error)
 }
 
 var (
@@ -79,7 +78,11 @@ func decodeUnsafe(bs []byte) []float32 {
 	return unsafe.Slice((*float32)(unsafe.Pointer(&bs[0])), len(bs)/4)
 }
 
-func (p ProviderSQL) SearchClosest(q []float32) (*models.VectorRow, error) {
+func (p ProviderSQL) SearchClosest(q []float32) ([]models.VectorRow, error) {
+	tableName, err := fetchTableName(q)
+	if err != nil {
+		return nil, err
+	}
 	stmt, _, err := p.s3Conn.Prepare(
 		fmt.Sprintf(`SELECT
 			id,
@@ -91,35 +94,35 @@ func (p ProviderSQL) SearchClosest(q []float32) (*models.VectorRow, error) {
 		WHERE embedding MATCH ?
 		ORDER BY distance
 		LIMIT 4
-	`, vecTableName))
+	`, tableName))
 	if err != nil {
-		log.Fatal(err)
+		return nil, err
 	}
 	query, err := sqlite_vec.SerializeFloat32(q[:])
 	if err != nil {
-		log.Fatal(err)
+		return nil, err
 	}
 	if err := stmt.BindBlob(1, query); err != nil {
 		p.logger.Error("failed to bind", "error", err)
 		return nil, err
 	}
-	resp := make([]models.VectorRow, 4)
-	i := 0
+	resp := []models.VectorRow{}
 	for stmt.Step() {
-		resp[i].ID = uint32(stmt.ColumnInt64(0))
-		resp[i].Distance = float32(stmt.ColumnFloat(1))
+		res := models.VectorRow{}
+		res.ID = uint32(stmt.ColumnInt64(0))
+		res.Distance = float32(stmt.ColumnFloat(1))
 		emb := stmt.ColumnRawText(2)
-		resp[i].Embeddings = decodeUnsafe(emb)
-		resp[i].Slug = stmt.ColumnText(3)
-		resp[i].RawText = stmt.ColumnText(4)
-		i++
+		res.Embeddings = decodeUnsafe(emb)
+		res.Slug = stmt.ColumnText(3)
+		res.RawText = stmt.ColumnText(4)
+		resp = append(resp, res)
 	}
 	if err := stmt.Err(); err != nil {
-		log.Fatal(err)
+		return nil, err
 	}
 	err = stmt.Close()
 	if err != nil {
-		log.Fatal(err)
+		return nil, err
 	}
-	return nil, nil
+	return resp, nil
 }

tui.go

@@ -5,6 +5,7 @@ import (
 	"elefant/pngmeta"
 	"fmt"
 	"os"
+	"path"
 	"strconv"
 	"strings"
 	"time"
@@ -169,8 +170,9 @@ func makeRAGTable(fileList []string) *tview.Table {
 		// notification := fmt.Sprintf("chat: %s; action: %s", fpath, tc.Text)
 		switch tc.Text {
 		case "load":
-			if err := loadRAG(fpath); err != nil {
-				logger.Error("failed to read history file", "chat", fpath)
+			fpath = path.Join(cfg.RAGDir, fpath)
+			if err := ragger.LoadRAG(fpath); err != nil {
+				logger.Error("failed to embed file", "chat", fpath, "error", err)
 				pages.RemovePage(RAGPage)
 				return
 			}
@@ -228,7 +230,7 @@ func colorText() {
 }
 
 func updateStatusLine() {
-	position.SetText(fmt.Sprintf(indexLine, botRespMode, cfg.AssistantRole, activeChatName, cfg.RAGEnabled))
+	position.SetText(fmt.Sprintf(indexLine, botRespMode, cfg.AssistantRole, activeChatName, cfg.RAGEnabled, cfg.EmbedURL))
 }
 
 func initSysCards() ([]string, error) {