import { text, tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./configSchematics";

function stripSearchMarkup(text: string): string {
  // Удаляем разметку поиска, которую вставляет API Wikipedia.
  return text.replaceAll('<span class="searchmatch">', "").replaceAll("</span>", "");
}

export async function toolsProvider(ctl: ToolsProviderController) {
  const config = ctl.getPluginConfig(configSchematics);

  const searchWikipediaTool = tool({
    name: "search_wikipedia_ru",
    description: text`
      Ищет информацию в Wikipedia по заданной строке \`query\`. Возвращает список результатов поиска.
      Каждый результат содержит \`title\` (заголовок), \`summary\` (краткий отрывок) и \`page_id\`, который можно использовать 
      для получения полного содержимого страницы с помощью инструмента get_wikipedia_page_ru.

      Примечание: этот инструмент работает с Wikipedia, поэтому вместо естественных вопросов 
      используйте ключевые слова или названия статей. Например, если пользователь спрашивает 
      "кто такой Альберт Эйнштейн", ищите "Альберт Эйнштейн", а не "кто такой Альберт Эйнштейн".

      Если конкретный запрос не дал ожидаемого результата, попробуйте использовать более 
      каноничное название или другой термин, который может содержать нужную информацию.

      ВСЕГДА используйте \`get_wikipedia_page_ru\` для получения полного содержания страницы после поиска. 
      НИКОГДА не пытайтесь дать окончательный ответ, основываясь только на кратком отрывке (summary) из результатов поиска.
    `,
    parameters: { query: z.string() },
    implementation: async ({ query }, { warn }) => {
      const url = new URL(config.get("wikipediaBaseUrl"));
      url.pathname = "/w/api.php";
      const searchParams = url.searchParams;
      searchParams.set("action", "query");
      searchParams.set("list", "search");
      searchParams.set("srsearch", query);
      searchParams.set("format", "json");
      searchParams.set("utf8", "1");

      const response = await fetch(url.toString());
      if (!response.ok) {
        return "Ошибка: Не удалось получить результаты поиска из Wikipedia.";
      }
      const data = await response.json();

      if (data.error !== undefined) {
        warn(`Wikipedia API returned an error: ${data.error.info}`);
        return data.error;
      }

      return {
        results: data.query.search.map(
          (result: { title: string; snippet: string; pageid: number }) => ({
            title: result.title,
            summary: stripSearchMarkup(result.snippet),
            page_id: result.pageid,
          }),
        ),
        hint: text`
          Если какой-либо из результатов поиска кажется подходящим, ОБЯЗАТЕЛЬНО используйте \`get_wikipedia_page_ru\`, 
          чтобы получить полное содержимое страницы по \`page_id\`. Поле \`summary\` — это лишь краткий фрагмент, 
          в котором может отсутствовать важная информация. Если подходящих результатов нет, попробуйте изменить 
          запрос, используя более точный термин.
        `,
      };
    },
  });

  const getWikipediaPageTool = tool({
    name: "get_wikipedia_page_ru",
    description: text`
      Извлекает полное текстовое содержимое страницы Wikipedia по её \`page_id\`. 
      Возвращает заголовок и текст страницы. Сначала используйте \`search_wikipedia\`, чтобы получить \`page_id\`.
    `,
    parameters: { page_id: z.number() },
    implementation: async ({ page_id }) => {
      const url = new URL(config.get("wikipediaBaseUrl"));
      url.pathname = "/w/api.php";
      const searchParams = url.searchParams;
      searchParams.set("action", "query");
      searchParams.set("prop", "extracts");
      searchParams.set("explaintext", "1");
      searchParams.set("pageids", String(page_id));
      searchParams.set("format", "json");
      searchParams.set("utf8", "1");

      const response = await fetch(url.toString());
      if (!response.ok) {
        return "Ошибка: Не удалось получить содержимое страницы из Wikipedia.";
      }
      const data = await response.json();
      const page = data.query.pages[page_id];

      return {
        title: page.title,
        content: page.extract ?? "Содержимое для этой страницы недоступно.",
      };
    },
  });

  return [searchWikipediaTool, getWikipediaPageTool];
}

import { text, tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./configSchematics";

function stripSearchMarkup(text: string): string {
  // Удаляем разметку поиска, которую вставляет API Wikipedia.
  return text.replaceAll('<span class="searchmatch">', "").replaceAll("</span>", "");
}

export async function toolsProvider(ctl: ToolsProviderController) {
  const config = ctl.getPluginConfig(configSchematics);

  const searchWikipediaTool = tool({
    name: "search_wikipedia_ru",
    description: text`
      Ищет информацию в Wikipedia по заданной строке \`query\`. Возвращает список результатов поиска.
      Каждый результат содержит \`title\` (заголовок), \`summary\` (краткий отрывок) и \`page_id\`, который можно использовать 
      для получения полного содержимого страницы с помощью инструмента get_wikipedia_page_ru.

      Примечание: этот инструмент работает с Wikipedia, поэтому вместо естественных вопросов 
      используйте ключевые слова или названия статей. Например, если пользователь спрашивает 
      "кто такой Альберт Эйнштейн", ищите "Альберт Эйнштейн", а не "кто такой Альберт Эйнштейн".

      Если конкретный запрос не дал ожидаемого результата, попробуйте использовать более 
      каноничное название или другой термин, который может содержать нужную информацию.

      ВСЕГДА используйте \`get_wikipedia_page_ru\` для получения полного содержания страницы после поиска. 
      НИКОГДА не пытайтесь дать окончательный ответ, основываясь только на кратком отрывке (summary) из результатов поиска.
    `,
    parameters: { query: z.string() },
    implementation: async ({ query }, { warn }) => {
      const url = new URL(config.get("wikipediaBaseUrl"));
      url.pathname = "/w/api.php";
      const searchParams = url.searchParams;
      searchParams.set("action", "query");
      searchParams.set("list", "search");
      searchParams.set("srsearch", query);
      searchParams.set("format", "json");
      searchParams.set("utf8", "1");

      const response = await fetch(url.toString());
      if (!response.ok) {
        return "Ошибка: Не удалось получить результаты поиска из Wikipedia.";
      }
      const data = await response.json();

      if (data.error !== undefined) {
        warn(`Wikipedia API returned an error: ${data.error.info}`);
        return data.error;
      }

      return {
        results: data.query.search.map(
          (result: { title: string; snippet: string; pageid: number }) => ({
            title: result.title,
            summary: stripSearchMarkup(result.snippet),
            page_id: result.pageid,
          }),
        ),
        hint: text`
          Если какой-либо из результатов поиска кажется подходящим, ОБЯЗАТЕЛЬНО используйте \`get_wikipedia_page_ru\`, 
          чтобы получить полное содержимое страницы по \`page_id\`. Поле \`summary\` — это лишь краткий фрагмент, 
          в котором может отсутствовать важная информация. Если подходящих результатов нет, попробуйте изменить 
          запрос, используя более точный термин.
        `,
      };
    },
  });

  const getWikipediaPageTool = tool({
    name: "get_wikipedia_page_ru",
    description: text`
      Извлекает полное текстовое содержимое страницы Wikipedia по её \`page_id\`. 
      Возвращает заголовок и текст страницы. Сначала используйте \`search_wikipedia\`, чтобы получить \`page_id\`.
    `,
    parameters: { page_id: z.number() },
    implementation: async ({ page_id }) => {
      const url = new URL(config.get("wikipediaBaseUrl"));
      url.pathname = "/w/api.php";
      const searchParams = url.searchParams;
      searchParams.set("action", "query");
      searchParams.set("prop", "extracts");
      searchParams.set("explaintext", "1");
      searchParams.set("pageids", String(page_id));
      searchParams.set("format", "json");
      searchParams.set("utf8", "1");

      const response = await fetch(url.toString());
      if (!response.ok) {
        return "Ошибка: Не удалось получить содержимое страницы из Wikipedia.";
      }
      const data = await response.json();
      const page = data.query.pages[page_id];

      return {
        title: page.title,
        content: page.extract ?? "Содержимое для этой страницы недоступно.",
      };
    },
  });

  return [searchWikipediaTool, getWikipediaPageTool];
}