//Copyright 2024 Google LLC. under the Apache License, Version 2.0
import "jsr:@std/dotenv@0.225/load";
import { GoogleGenerativeAI } from "npm:@google/generative-ai@0.21.0";
import { sumOf } from "jsr:@std/collections@1/sum-of";

const texts = [
  "The quick brown fox jumps over the lazy dog.",
  "The quick rbown fox jumps over the lazy dog.",
  "teh fast fox jumps over the slow woofer.",
  "a quick brown fox jmps over lazy dog.",
  "brown fox jumping over dog",
  "fox > dog",
  // Alternative pangram for comparison:
  "The five boxing wizards jump quickly.",
  // Unrelated text, also for comparison:
  "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Vivamus et hendrerit massa. Sed pulvinar, nisi a lobortis sagittis, neque risus gravida dolor, in porta dui odio vel purus.",
];

const ai = new GoogleGenerativeAI(Deno.env.get("GOOGLE_API_KEY")!);
const model = ai.getGenerativeModel({ model: "models/text-embedding-004" });
const response = await model.batchEmbedContents({
  requests: texts.map((text) => ({
    content: { role: "user", parts: [{ text }] },
  })),
});

const matrix = response.embeddings.map((embedding) => embedding.values);
const similarities = matrix.map((row) =>
  matrix.map((col) => sumOf(row.map((v, j) => v * col[j]), (i) => i))
);

console.table(similarities.map((row) => row.map((v) => v.toPrecision(3))));