diff options
author | Grail Finder <wohilas@gmail.com> | 2024-12-31 13:25:13 +0300 |
---|---|---|
committer | Grail Finder <wohilas@gmail.com> | 2024-12-31 13:25:13 +0300 |
commit | 461d19aa2512fea7ac07e50c3178609850ef07c3 (patch) | |
tree | 1e317b813686ca7696e96bba766cdbd3aeb35138 | |
parent | 4db8aea43dded3f2c1d0d41b5a3fb322a38d4730 (diff) |
Feat: add rag [wip; skip-ci]
-rw-r--r-- | README.md | 1 | ||||
-rw-r--r-- | bot.go | 3 | ||||
-rw-r--r-- | config/config.go | 2 | ||||
-rw-r--r-- | go.mod | 10 | ||||
-rw-r--r-- | go.sum | 18 | ||||
-rw-r--r-- | models/db.go | 10 | ||||
-rw-r--r-- | models/models.go | 19 | ||||
-rw-r--r-- | rag.go | 40 | ||||
-rw-r--r-- | storage/migrate.go | 5 | ||||
-rw-r--r-- | storage/migrations/002_add_vector.up.sql | 6 | ||||
-rw-r--r-- | storage/storage.go | 9 | ||||
-rw-r--r-- | storage/storage_test.go | 88 | ||||
-rw-r--r-- | storage/vector.go | 89 |
13 files changed, 292 insertions, 8 deletions
@@ -52,3 +52,4 @@ - user icon (and role?) from config is not used; + - message editing broke ( runtime error: index out of range [-1]); + - F1 can load any chat, by loading chat of other agent it does not switch agents, if that chat is continued, it will rewrite agent in db; (either allow only chats from current agent OR switch agent on chat loading); +- after chat is deleted: load undeleted chat; @@ -278,6 +278,9 @@ func init() { // logger = slog.New(slog.NewTextHandler(logfile, nil)) store = storage.NewProviderSQL("test.db", logger) + if store == nil { + os.Exit(1) + } // https://github.com/coreydaley/ggerganov-llama.cpp/blob/master/examples/server/README.md // load all chats in memory if _, err := loadHistoryChats(); err != nil { diff --git a/config/config.go b/config/config.go index 074bfda..ce1b877 100644 --- a/config/config.go +++ b/config/config.go @@ -8,6 +8,7 @@ import ( type Config struct { APIURL string `toml:"APIURL"` + EmbedURL string `toml:"EmbedURL"` ShowSys bool `toml:"ShowSys"` LogFile string `toml:"LogFile"` UserRole string `toml:"UserRole"` @@ -29,6 +30,7 @@ func LoadConfigOrDefault(fn string) *Config { if err != nil { fmt.Println("failed to read config from file, loading default") config.APIURL = "http://localhost:8080/v1/chat/completions" + config.EmbedURL = "http://localhost:8080/v1/embiddings" config.ShowSys = true config.LogFile = "log.txt" config.UserRole = "user" @@ -4,24 +4,28 @@ go 1.23.2 require ( github.com/BurntSushi/toml v1.4.0 + github.com/asg017/sqlite-vec-go-bindings v0.1.6 github.com/gdamore/tcell/v2 v2.7.4 github.com/glebarez/go-sqlite v1.22.0 github.com/jmoiron/sqlx v1.4.0 + github.com/ncruces/go-sqlite3 v0.21.3 github.com/rivo/tview v0.0.0-20241103174730-c76f7879f592 ) require ( github.com/dustin/go-humanize v1.0.1 // indirect github.com/gdamore/encoding v1.0.0 // indirect - github.com/google/uuid v1.5.0 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/lucasb-eyer/go-colorful v1.2.0 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-runewidth v0.0.15 // indirect + github.com/ncruces/julianday v1.0.0 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/rivo/uniseg v0.4.7 // indirect - golang.org/x/sys v0.17.0 // indirect + github.com/tetratelabs/wazero v1.8.2 // indirect + golang.org/x/sys v0.28.0 // indirect golang.org/x/term v0.17.0 // indirect - golang.org/x/text v0.14.0 // indirect + golang.org/x/text v0.21.0 // indirect modernc.org/libc v1.37.6 // indirect modernc.org/mathutil v1.6.0 // indirect modernc.org/memory v1.7.2 // indirect @@ -2,6 +2,8 @@ filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA= filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= github.com/BurntSushi/toml v1.4.0 h1:kuoIxZQy2WRRk1pttg9asf+WVv6tWQuBNVmK8+nqPr0= github.com/BurntSushi/toml v1.4.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= +github.com/asg017/sqlite-vec-go-bindings v0.1.6 h1:Nx0jAzyS38XpkKznJ9xQjFXz2X9tI7KqjwVxV8RNoww= +github.com/asg017/sqlite-vec-go-bindings v0.1.6/go.mod h1:A8+cTt/nKFsYCQF6OgzSNpKZrzNo5gQsXBTfsXHXY0Q= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/gdamore/encoding v1.0.0 h1:+7OoQ1Bc6eTm5niUzBa0Ctsh6JbMW6Ra+YNuAtDBdko= @@ -14,8 +16,8 @@ github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpv github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg= github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26 h1:Xim43kblpZXfIBQsbuBVKCudVG457BR2GZFIz3uw3hQ= github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo= -github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU= -github.com/google/uuid v1.5.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o= github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= @@ -28,6 +30,10 @@ github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZ github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= +github.com/ncruces/go-sqlite3 v0.21.3 h1:hHkfNQLcbnxPJZhC/RGw9SwP3bfkv/Y0xUHWsr1CdMQ= +github.com/ncruces/go-sqlite3 v0.21.3/go.mod h1:zxMOaSG5kFYVFK4xQa0pdwIszqxqJ0W0BxBgwdrNjuA= +github.com/ncruces/julianday v1.0.0 h1:fH0OKwa7NWvniGQtxdJRxAgkBMolni2BjDHaWTxqt7M= +github.com/ncruces/julianday v1.0.0/go.mod h1:Dusn2KvZrrovOMJuOt0TNXL6tB7U2E8kvza5fFc9G7g= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rivo/tview v0.0.0-20241103174730-c76f7879f592 h1:YIJ+B1hePP6AgynC5TcqpO0H9k3SSoZa2BGyL6vDUzM= @@ -36,6 +42,8 @@ github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJ github.com/rivo/uniseg v0.4.3/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/tetratelabs/wazero v1.8.2 h1:yIgLR/b2bN31bjxwXHD8a3d+BogigR952csSDdLYEv4= +github.com/tetratelabs/wazero v1.8.2/go.mod h1:yAI0XTsMBhREkM/YDAK/zNou3GoiAce1P6+rp/wQhjs= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= @@ -55,8 +63,9 @@ golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y= golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= +golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= @@ -66,8 +75,9 @@ 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/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= -golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= +golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= +golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= diff --git a/models/db.go b/models/db.go index 1f1c8b4..db4ae75 100644 --- a/models/db.go +++ b/models/db.go @@ -35,3 +35,13 @@ type Memory struct { CreatedAt time.Time `db:"created_at" json:"created_at"` UpdatedAt time.Time `db:"updated_at" json:"updated_at"` } + +// vector models + +type VectorRow struct { + ID uint32 `db:"id" json:"id"` + Embeddings []float32 `db:"embeddings" json:"embeddings"` + Slug string `db:"slug" json:"slug"` + RawText string `db:"raw_text" json:"raw_text"` + Distance float32 `db:"distance" json:"distance"` +} diff --git a/models/models.go b/models/models.go index 2c5c99c..bbf1bd0 100644 --- a/models/models.go +++ b/models/models.go @@ -106,3 +106,22 @@ type ChatToolsBody struct { } `json:"tools"` ToolChoice string `json:"tool_choice"` } + +type EmbeddingResp struct { + Embedding []float32 `json:"embedding"` + Index uint32 `json:"index"` +} + +// type EmbeddingsResp struct { +// Model string `json:"model"` +// Object string `json:"object"` +// Usage struct { +// PromptTokens int `json:"prompt_tokens"` +// TotalTokens int `json:"total_tokens"` +// } `json:"usage"` +// Data []struct { +// Embedding []float32 `json:"embedding"` +// Index int `json:"index"` +// Object string `json:"object"` +// } `json:"data"` +// } @@ -0,0 +1,40 @@ +package main + +import ( + "bytes" + "elefant/models" + "encoding/json" +) + +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 + } + 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() + 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 + } + 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([5120]float32(emb.Embedding)) +} diff --git a/storage/migrate.go b/storage/migrate.go index d97b99d..3c4ca21 100644 --- a/storage/migrate.go +++ b/storage/migrate.go @@ -5,6 +5,8 @@ import ( "fmt" "io/fs" "strings" + + _ "github.com/asg017/sqlite-vec-go-bindings/ncruces" ) //go:embed migrations/* @@ -27,6 +29,7 @@ func (p *ProviderSQL) Migrate() { err := p.executeMigration(migrationsDir, file.Name()) if err != nil { p.logger.Error("Failed to execute migration %s: %v", file.Name(), err) + panic(err) } } } @@ -51,7 +54,7 @@ func (p *ProviderSQL) executeMigration(migrationsDir fs.FS, fileName string) err func (p *ProviderSQL) executeSQL(sqlContent []byte) error { // Connect to the database (example using a simple connection) - _, err := p.db.Exec(string(sqlContent)) + err := p.s3Conn.Exec(string(sqlContent)) if err != nil { return fmt.Errorf("failed to execute SQL: %w", err) } diff --git a/storage/migrations/002_add_vector.up.sql b/storage/migrations/002_add_vector.up.sql new file mode 100644 index 0000000..4fcc9aa --- /dev/null +++ b/storage/migrations/002_add_vector.up.sql @@ -0,0 +1,6 @@ +CREATE VIRTUAL TABLE IF NOT EXISTS embeddings USING vec0( + id INTEGER PRIMARY KEY AUTOINCREMENT, + embedding FLOAT[5120], + slug TEXT NOT NULL, + raw_text TEXT NOT NULL +); diff --git a/storage/storage.go b/storage/storage.go index 0853328..cfaaf0a 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -6,11 +6,13 @@ import ( _ "github.com/glebarez/go-sqlite" "github.com/jmoiron/sqlx" + "github.com/ncruces/go-sqlite3" ) type FullRepo interface { ChatHistory Memories + VectorRepo } type ChatHistory interface { @@ -25,6 +27,7 @@ type ChatHistory interface { type ProviderSQL struct { db *sqlx.DB + s3Conn *sqlite3.Conn logger *slog.Logger } @@ -87,6 +90,7 @@ func (p ProviderSQL) ChatGetMaxID() (uint32, error) { return id, err } +// opens two connections func NewProviderSQL(dbPath string, logger *slog.Logger) FullRepo { db, err := sqlx.Open("sqlite", dbPath) if err != nil { @@ -94,6 +98,11 @@ func NewProviderSQL(dbPath string, logger *slog.Logger) FullRepo { return nil } p := ProviderSQL{db: db, logger: logger} + p.s3Conn, err = sqlite3.Open(dbPath) + if err != nil { + logger.Error("failed to open vecdb connection", "error", err) + return nil + } p.Migrate() return p } diff --git a/storage/storage_test.go b/storage/storage_test.go index 8373ab0..f6af4f5 100644 --- a/storage/storage_test.go +++ b/storage/storage_test.go @@ -3,13 +3,16 @@ package storage import ( "elefant/models" "fmt" + "log" "log/slog" "os" "testing" "time" + sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/ncruces" _ "github.com/glebarez/go-sqlite" "github.com/jmoiron/sqlx" + "github.com/ncruces/go-sqlite3" ) func TestMemories(t *testing.T) { @@ -160,3 +163,88 @@ func TestChatHistory(t *testing.T) { t.Errorf("Expected 0 chats, got %d", len(chats)) } } + +func TestVecTable(t *testing.T) { + // healthcheck + db, err := sqlite3.Open(":memory:") + if err != nil { + t.Fatal(err) + } + stmt, _, err := db.Prepare(`SELECT sqlite_version(), vec_version()`) + if err != nil { + t.Fatal(err) + } + stmt.Step() + log.Printf("sqlite_version=%s, vec_version=%s\n", stmt.ColumnText(0), stmt.ColumnText(1)) + stmt.Close() + // migration + err = db.Exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4], chat_name TEXT NOT NULL)") + if err != nil { + t.Fatal(err) + } + // data prep and insert + items := map[int][]float32{ + 1: {0.1, 0.1, 0.1, 0.1}, + 2: {0.2, 0.2, 0.2, 0.2}, + 3: {0.3, 0.3, 0.3, 0.3}, + 4: {0.4, 0.4, 0.4, 0.4}, + 5: {0.5, 0.5, 0.5, 0.5}, + } + q := []float32{0.28, 0.3, 0.3, 0.3} + stmt, _, err = db.Prepare("INSERT INTO vec_items(rowid, embedding, chat_name) VALUES (?, ?, ?)") + if err != nil { + t.Fatal(err) + } + for id, values := range items { + v, err := sqlite_vec.SerializeFloat32(values) + if err != nil { + t.Fatal(err) + } + stmt.BindInt(1, id) + stmt.BindBlob(2, v) + stmt.BindText(3, "some_chat") + err = stmt.Exec() + if err != nil { + t.Fatal(err) + } + stmt.Reset() + } + stmt.Close() + // select | vec search + stmt, _, err = db.Prepare(` + SELECT + rowid, + distance, + embedding + FROM vec_items + WHERE embedding MATCH ? + ORDER BY distance + LIMIT 3 + `) + if err != nil { + t.Fatal(err) + } + query, err := sqlite_vec.SerializeFloat32(q) + if err != nil { + t.Fatal(err) + } + stmt.BindBlob(1, query) + for stmt.Step() { + rowid := stmt.ColumnInt64(0) + distance := stmt.ColumnFloat(1) + emb := stmt.ColumnRawText(2) + floats := decodeUnsafe(emb) + log.Printf("rowid=%d, distance=%f, floats=%v\n", rowid, distance, floats) + } + if err := stmt.Err(); err != nil { + t.Fatal(err) + } + err = stmt.Close() + if err != nil { + t.Fatal(err) + } + err = db.Close() + if err != nil { + t.Fatal(err) + } +} diff --git a/storage/vector.go b/storage/vector.go new file mode 100644 index 0000000..bc46734 --- /dev/null +++ b/storage/vector.go @@ -0,0 +1,89 @@ +package storage + +import ( + "elefant/models" + "fmt" + "log" + "unsafe" + + sqlite_vec "github.com/asg017/sqlite-vec-go-bindings/ncruces" +) + +type VectorRepo interface { + WriteVector(*models.VectorRow) error + SearchClosest(q [5120]float32) (*models.VectorRow, error) +} + +var vecTableName = "embeddings" + +func (p ProviderSQL) WriteVector(row *models.VectorRow) error { + stmt, _, err := p.s3Conn.Prepare( + fmt.Sprintf("INSERT INTO %s(embedding, slug, raw_text) VALUES (?, ?, ?)", vecTableName)) + defer stmt.Close() + if err != nil { + p.logger.Error("failed to prep a stmt", "error", err) + return err + } + 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) + err = stmt.Exec() + if err != nil { + p.logger.Error("failed exec a stmt", "error", err) + return err + } + return nil +} + +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 + id, + distance, + embedding, + slug, + raw_text + FROM vec_items + WHERE embedding MATCH ? + ORDER BY distance + LIMIT 4 + `) + if err != nil { + log.Fatal(err) + } + query, err := sqlite_vec.SerializeFloat32(q[:]) + if err != nil { + log.Fatal(err) + } + stmt.BindBlob(1, query) + resp := make([]models.VectorRow, 4) + i := 0 + for stmt.Step() { + resp[i].ID = uint32(stmt.ColumnInt64(0)) + resp[i].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++ + } + if err := stmt.Err(); err != nil { + log.Fatal(err) + } + err = stmt.Close() + if err != nil { + log.Fatal(err) + } + return nil, nil +} |