Radu Macocian (admac)
48 lines
1.2 KiB
Go
48 lines
1.2 KiB
Go
package embed
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
|
|
"git.estatecloud.org/radumaco/souvenir/config"
|
|
"github.com/gopsql/pgx"
|
|
)
|
|
|
|
type Embed struct {
|
|
Cfg config.DbConfig
|
|
Logger slog.Logger
|
|
}
|
|
|
|
func (e Embed) tableName() string {
|
|
return fmt.Sprintf(`message_chunks_vec_%s`, e.Cfg.Embedding.Model)
|
|
}
|
|
|
|
func (e Embed) Init() error {
|
|
const createVectorTableTemplate = `
|
|
CREATE TABLE IF NOT EXISTS %s (
|
|
chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE,
|
|
embedding vector(384),
|
|
content_hash text not null
|
|
)`
|
|
conn := pgx.MustOpen(e.Cfg.ConnectionString())
|
|
var exists bool
|
|
if err := conn.QueryRow(`SELECT EXISTS
|
|
(SELECT 1 FROM information_schema.tables
|
|
WHERE table_schema = 'public'
|
|
AND table_name =$1)`, e.tableName()).Scan(&exists); err != nil {
|
|
e.Logger.Error("Error while checking table", "table", e.tableName(), "error", err)
|
|
return err
|
|
}
|
|
if exists {
|
|
e.Logger.Debug("Table found", "table", e.tableName())
|
|
} else {
|
|
e.Logger.Debug("Table not found. Creating", "table", e.tableName())
|
|
script := fmt.Sprintf(createVectorTableTemplate, e.tableName())
|
|
if _, err := conn.Exec(script); err != nil {
|
|
e.Logger.Error("Could not create table", "table", e.tableName(), "error", err)
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|