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

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

  const searchWebTool = tool({
    name: "search_web",
    description: text`
      Searches the web using the given \`query\` string. Returns a list of results, each with a
      \`title\`, \`url\`, \`snippet\`, \`source\` (the site's hostname), and \`published_date\` (if
      known, otherwise null).

      Use natural search terms, similar to what you'd type into a search engine, rather than full
      natural language questions.

      ALWAYS use \`fetch_page\` on a promising result's \`url\` to read its full content before
      answering. NEVER answer merely based on the \`snippet\`, which is often incomplete or
      misleading.
    `,
    parameters: {
      query: z.string(),
      count: z.number().int().min(1).max(20).optional(),
    },
    implementation: async ({ query, count }, { warn }) => {
      const timeoutMs = config.get("fetchTimeoutMs");
      try {
        const outcome = await searchWeb({
          searxngBaseUrl: config.get("searxngBaseUrl"),
          braveBaseUrl: config.get("braveBaseUrl"),
          braveApiKey: config.get("braveApiKey"),
          query,
          count: count ?? config.get("defaultResultCount"),
          timeoutMs,
        });

        for (const warning of outcome.warnings) {
          warn(warning);
        }

        return {
          provider: outcome.provider,
          results: outcome.results,
          hint: text`
            If any of the results look relevant, use \`fetch_page\` on their \`url\` to read the
            full content. If none look relevant, try again with different search terms.
          `,
        };
      } catch (error) {
        return `Error: Failed to search the web. ${(error as Error).message}`;
      }
    },
  });

  const fetchPageTool = tool({
    name: "fetch_page",
    description: text`
      Fetches a web page at the given \`url\` and returns its main content converted to Markdown,
      with navigation, ads, and scripts stripped out. Also returns citation metadata: \`title\`,
      \`url\`, \`source\`, and \`published_date\` (if known).

      Long pages are returned in chunks. The response includes a \`chunk\` object with
      \`has_more\` and \`next_offset\`. If \`has_more\` is true and you need more of the page, call
      \`fetch_page\` again with the same \`url\` and \`offset\` set to \`next_offset\`.
    `,
    parameters: {
      url: z.string(),
      offset: z.number().int().min(0).optional(),
    },
    implementation: async ({ url, offset }) => {
      const timeoutMs = config.get("fetchTimeoutMs");
      const chunkSize = config.get("fetchChunkSize");

      let page;
      try {
        page = await fetchAndExtract(url, timeoutMs);
      } catch (error) {
        return `Error: Failed to fetch or extract page content. ${(error as Error).message}`;
      }

      const start = offset ?? 0;
      const end = Math.min(start + chunkSize, page.content.length);
      const hasMore = end < page.content.length;

      return {
        title: page.title,
        url: page.url,
        source: page.source,
        published_date: page.published_date,
        content: page.content.slice(start, end),
        chunk: {
          offset: start,
          next_offset: hasMore ? end : null,
          total_length: page.content.length,
          has_more: hasMore,
        },
        ...(hasMore
          ? {
              hint: text`
                This page has more content. Call \`fetch_page\` again with \`url\` set to
                "${page.url}" and \`offset\` set to ${end} to continue reading.
              `,
            }
          : {}),
      };
    },
  });

  return [searchWebTool, fetchPageTool];
}

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

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

  const searchWebTool = tool({
    name: "search_web",
    description: text`
      Searches the web using the given \`query\` string. Returns a list of results, each with a
      \`title\`, \`url\`, \`snippet\`, \`source\` (the site's hostname), and \`published_date\` (if
      known, otherwise null).

      Use natural search terms, similar to what you'd type into a search engine, rather than full
      natural language questions.

      ALWAYS use \`fetch_page\` on a promising result's \`url\` to read its full content before
      answering. NEVER answer merely based on the \`snippet\`, which is often incomplete or
      misleading.
    `,
    parameters: {
      query: z.string(),
      count: z.number().int().min(1).max(20).optional(),
    },
    implementation: async ({ query, count }, { warn }) => {
      const timeoutMs = config.get("fetchTimeoutMs");
      try {
        const outcome = await searchWeb({
          searxngBaseUrl: config.get("searxngBaseUrl"),
          braveBaseUrl: config.get("braveBaseUrl"),
          braveApiKey: config.get("braveApiKey"),
          query,
          count: count ?? config.get("defaultResultCount"),
          timeoutMs,
        });

        for (const warning of outcome.warnings) {
          warn(warning);
        }

        return {
          provider: outcome.provider,
          results: outcome.results,
          hint: text`
            If any of the results look relevant, use \`fetch_page\` on their \`url\` to read the
            full content. If none look relevant, try again with different search terms.
          `,
        };
      } catch (error) {
        return `Error: Failed to search the web. ${(error as Error).message}`;
      }
    },
  });

  const fetchPageTool = tool({
    name: "fetch_page",
    description: text`
      Fetches a web page at the given \`url\` and returns its main content converted to Markdown,
      with navigation, ads, and scripts stripped out. Also returns citation metadata: \`title\`,
      \`url\`, \`source\`, and \`published_date\` (if known).

      Long pages are returned in chunks. The response includes a \`chunk\` object with
      \`has_more\` and \`next_offset\`. If \`has_more\` is true and you need more of the page, call
      \`fetch_page\` again with the same \`url\` and \`offset\` set to \`next_offset\`.
    `,
    parameters: {
      url: z.string(),
      offset: z.number().int().min(0).optional(),
    },
    implementation: async ({ url, offset }) => {
      const timeoutMs = config.get("fetchTimeoutMs");
      const chunkSize = config.get("fetchChunkSize");

      let page;
      try {
        page = await fetchAndExtract(url, timeoutMs);
      } catch (error) {
        return `Error: Failed to fetch or extract page content. ${(error as Error).message}`;
      }

      const start = offset ?? 0;
      const end = Math.min(start + chunkSize, page.content.length);
      const hasMore = end < page.content.length;

      return {
        title: page.title,
        url: page.url,
        source: page.source,
        published_date: page.published_date,
        content: page.content.slice(start, end),
        chunk: {
          offset: start,
          next_offset: hasMore ? end : null,
          total_length: page.content.length,
          has_more: hasMore,
        },
        ...(hasMore
          ? {
              hint: text`
                This page has more content. Call \`fetch_page\` again with \`url\` set to
                "${page.url}" and \`offset\` set to ${end} to continue reading.
              `,
            }
          : {}),
      };
    },
  });

  return [searchWebTool, fetchPageTool];
}