// @ts-check

// run in your browser or with `node internet.mjs` with the ALBERT_API_KEY set below
// see also: https://albert.api.etalab.gouv.fr/documentation

const API_KEY = process.env.ALBERT_API_KEY;

const API_URL = "https://albert.api.etalab.gouv.fr";
const LANGUAGE_MODEL = "AgentPublic/llama3-instruct-8b"; // see https://albert.api.etalab.gouv.fr/v1/models

/**
 * Call albert API
 * @param {{path: string, method?: "POST"|"GET", body: string}} param0
 * @returns {Promise<any>} Search API result
 */
const albertApi = ({ path, method = "POST", body }) =>
  fetch(`${API_URL}/v1${path}`, {
    method,
    headers: {
      Authorization: `Bearer ${API_KEY}`,
      "Content-Type": "application/json",
    },
    body,
  }).then((r) => r.json());

/**
 * Query some albert collection
 *
 * @param {{query: string, collectionId: string}} param0
 * @returns {Promise<string>}
 */
const askAlbert = async ({ query, collectionId }) => {
  const searchResult = await albertApi({
    path: "/search",
    body: JSON.stringify({ collections: [collectionId], k: 6, prompt: query }),
  });

  const prompt = `Réponds à la question suivante en markdown en te basant sur les documents ci-dessous : ${query}
  
# Documents :
  
${searchResult.data.map((c) => c.chunk.content)}`;

  const result = await albertApi({
    path: "/chat/completions",
    body: JSON.stringify({
      model: LANGUAGE_MODEL,
      stream: false,
      messages: [{ role: "user", content: prompt }],
    }),
  });

  const sources = searchResult.data.map(
    (c) => `${c.chunk.metadata.document_name}`
  );

  const sourcesList = sources.length
    ? "Sources:\n" +
      Array.from(new Set(sources))
        .sort()
        .map((s) => `- ${s}`)
        .join("\n")
    : "";

  return `${result.choices[0].message.content}\n\n${sourcesList}`;
};

const queries = [
  "Qui est Ada Lovelace ?",
  "C'est quoi la direction interministerielle du numérique ?",
  "Qui est premier ministre en France ?",
];

queries.forEach(async (q) => {
  const answer = await askAlbert({ query: q, collectionId: "internet" });
  console.log(`\n\n${q}\n> ${answer}`);
});