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.
This commit is contained in:
@@ -2,6 +2,7 @@ package storage
|
||||
|
||||
import (
|
||||
"elefant/models"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"unsafe"
|
||||
@@ -11,29 +12,61 @@ import (
|
||||
|
||||
type VectorRepo interface {
|
||||
WriteVector(*models.VectorRow) error
|
||||
SearchClosest(q [5120]float32) (*models.VectorRow, error)
|
||||
SearchClosest(q []float32) (*models.VectorRow, error)
|
||||
}
|
||||
|
||||
var vecTableName = "embeddings"
|
||||
var (
|
||||
vecTableName = "embeddings"
|
||||
vecTableName384 = "embeddings_384"
|
||||
)
|
||||
|
||||
func fetchTableName(emb []float32) (string, error) {
|
||||
switch len(emb) {
|
||||
case 5120:
|
||||
return vecTableName, nil
|
||||
case 384:
|
||||
return vecTableName384, nil
|
||||
default:
|
||||
return "", fmt.Errorf("no table for the size of %d", len(emb))
|
||||
}
|
||||
}
|
||||
|
||||
func (p ProviderSQL) WriteVector(row *models.VectorRow) error {
|
||||
tableName, err := fetchTableName(row.Embeddings)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
stmt, _, err := p.s3Conn.Prepare(
|
||||
fmt.Sprintf("INSERT INTO %s(embedding, slug, raw_text) VALUES (?, ?, ?)", vecTableName))
|
||||
defer stmt.Close()
|
||||
fmt.Sprintf("INSERT INTO %s(embedding, slug, raw_text) VALUES (?, ?, ?)", tableName))
|
||||
if err != nil {
|
||||
p.logger.Error("failed to prep a stmt", "error", err)
|
||||
return err
|
||||
}
|
||||
defer stmt.Close()
|
||||
v, err := sqlite_vec.SerializeFloat32(row.Embeddings)
|
||||
if err != nil {
|
||||
p.logger.Error("failed to serialize vector",
|
||||
"emb-len", len(row.Embeddings), "error", err)
|
||||
return err
|
||||
}
|
||||
stmt.BindInt(1, int(row.ID))
|
||||
stmt.BindBlob(2, v)
|
||||
stmt.BindText(3, row.Slug)
|
||||
stmt.BindText(4, row.RawText)
|
||||
if v == nil {
|
||||
err = errors.New("empty vector after serialization")
|
||||
p.logger.Error("empty vector after serialization",
|
||||
"emb-len", len(row.Embeddings), "text", row.RawText, "error", err)
|
||||
return err
|
||||
}
|
||||
if err := stmt.BindBlob(1, v); err != nil {
|
||||
p.logger.Error("failed to bind", "error", err)
|
||||
return err
|
||||
}
|
||||
if err := stmt.BindText(2, row.Slug); err != nil {
|
||||
p.logger.Error("failed to bind", "error", err)
|
||||
return err
|
||||
}
|
||||
if err := stmt.BindText(3, row.RawText); err != nil {
|
||||
p.logger.Error("failed to bind", "error", err)
|
||||
return err
|
||||
}
|
||||
err = stmt.Exec()
|
||||
if err != nil {
|
||||
p.logger.Error("failed exec a stmt", "error", err)
|
||||
@@ -46,19 +79,19 @@ func decodeUnsafe(bs []byte) []float32 {
|
||||
return unsafe.Slice((*float32)(unsafe.Pointer(&bs[0])), len(bs)/4)
|
||||
}
|
||||
|
||||
func (p ProviderSQL) SearchClosest(q [5120]float32) (*models.VectorRow, error) {
|
||||
stmt, _, err := p.s3Conn.Prepare(`
|
||||
SELECT
|
||||
func (p ProviderSQL) SearchClosest(q []float32) (*models.VectorRow, error) {
|
||||
stmt, _, err := p.s3Conn.Prepare(
|
||||
fmt.Sprintf(`SELECT
|
||||
id,
|
||||
distance,
|
||||
embedding,
|
||||
slug,
|
||||
raw_text
|
||||
FROM vec_items
|
||||
FROM %s
|
||||
WHERE embedding MATCH ?
|
||||
ORDER BY distance
|
||||
LIMIT 4
|
||||
`)
|
||||
`, vecTableName))
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
@@ -66,7 +99,10 @@ func (p ProviderSQL) SearchClosest(q [5120]float32) (*models.VectorRow, error) {
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
stmt.BindBlob(1, query)
|
||||
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
|
||||
for stmt.Step() {
|
||||
|
||||
Reference in New Issue
Block a user