feat: get similar images function

This commit is contained in:
2025-05-13 17:03:05 +01:00
parent b58efc413d
commit 85ff91a11b

View File

@ -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}
} }