embed.go

func GetCosineSimilarity(a, b []float32, idx int) (float32, error) {
	if len(a) != len(b) {
		return 0, errors.New(fmt.Sprintf("vectors must have the same length: %d != %d, idx: %d", len(a), len(b), idx))
	}

	var dotProduct float32
	for i := 0; i < len(a); i++ {
		dotProduct += a[i] * b[i]
	}
  // openai vectors are already normed
	return dotProduct, nil
}

func getTopKMostSimilarColumns(nlqEmbedding []float32, columnEmbeddings []ColumnAndEmbedding, k int) []ColumnSimilarity {
	similarities := make([]ColumnSimilarity, len(columnEmbeddings))
	for i, ce := range columnEmbeddings {
		similarity, err := GetCosineSimilarity(nlqEmbedding, ce.Embedding, i)
		if err != nil {
			log.Printf("Error calculating similarity: %v", err)
			continue
		}
		similarities[i] = ColumnSimilarity{Column: ce.Column, Similarity: similarity}
	}

	sort.Slice(similarities, func(i, j int) bool {
		return similarities[i].Similarity > similarities[j].Similarity
	})

	return similarities[:k]
}