feat: get similar images function
This commit is contained in:
@ -196,6 +196,44 @@ func (m ImageModel) UpdateEmbedding(ctx context.Context, imageId uuid.UUID, embe
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const SIMILARITY_THRESHOLD = 0.9
|
||||||
|
|
||||||
|
func (m ImageModel) GetSimilar(ctx context.Context, imageId uuid.UUID) ([]model.Image, error) {
|
||||||
|
stmt, err := m.dbPool.PrepareContext(ctx, `
|
||||||
|
WITH similarities AS (
|
||||||
|
SELECT 1 - (embedding <=> (
|
||||||
|
SELECT embedding
|
||||||
|
FROM haystack.image
|
||||||
|
WHERE ID = $1
|
||||||
|
)) AS cosine_similarity, id, image_name FROM haystack.image
|
||||||
|
WHERE embedding IS NOT NULL
|
||||||
|
)
|
||||||
|
SELECT id, image_name
|
||||||
|
FROM similatiries
|
||||||
|
WHERE cosine_similarity > $2
|
||||||
|
`)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return []model.Image{}, err
|
||||||
|
}
|
||||||
|
|
||||||
|
images := []model.Image{}
|
||||||
|
|
||||||
|
rows, err := stmt.QueryContext(ctx, imageId, SIMILARITY_THRESHOLD)
|
||||||
|
|
||||||
|
for rows.Next() {
|
||||||
|
image := model.Image{}
|
||||||
|
err := rows.Scan(&image)
|
||||||
|
if err != nil {
|
||||||
|
return []model.Image{}, err
|
||||||
|
}
|
||||||
|
|
||||||
|
images = append(images, image)
|
||||||
|
}
|
||||||
|
|
||||||
|
return images, nil
|
||||||
|
}
|
||||||
|
|
||||||
func NewImageModel(db *sql.DB) ImageModel {
|
func NewImageModel(db *sql.DB) ImageModel {
|
||||||
return ImageModel{dbPool: db}
|
return ImageModel{dbPool: db}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user