web-tools

/**
 * Fetch Images tool factory.
 */

import { tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk"
import { z } from "zod"

import { resolveConfig } from "../config/resolve-config"
import { formatToolError, UnsupportedContentTypeError } from "../errors"
import { createRetryNotifier, httpUrlSchema } from "../http"
import { downloadImages } from "../images"
import { fetchPage } from "../page"
import { extractPageImages, parseHtml } from "../parsers"
import { renderImageResults, type ImageSubject } from "../renderers"

import type { TTLCache } from "../cache"
import type { RetryOptions } from "../http"
import type { FetchedPage } from "../page"
import type { RateLimiter } from "../timing"
import type { Impit } from "impit"

/**
 * Build the Fetch Images tool, which downloads explicit image URLs or scrapes images from a given
 * page into the per-chat working directory and returns per-image markdown references.
 *
 * @param ctl - Tools provider controller supplied by the LM Studio SDK.
 * @param impit - Shared HTTP client used for HTML fetches and image downloads.
 * @param pageCache - Cache holding recent HTML payloads keyed by URL.
 * @param rateLimiter - Shared limiter enforcing the minimum gap between outbound requests.
 * @param imageLimiter - Shared limiter capping the number of image downloads in flight concurrently.
 * @param retry - Retry policy applied to every outbound request.
 * @returns The configured Fetch Images tool.
 */
export function createFetchImagesTool(
  ctl: ToolsProviderController,
  impit: Impit,
  pageCache: TTLCache<FetchedPage>,
  rateLimiter: RateLimiter,
  imageLimiter: RateLimiter,
  retry: RetryOptions
): Tool {
  return tool({
    name: "Fetch Images",
    description:
      "Download images to the chat's local working directory. Returns one record per image with a markdown reference (`![alt](localPath)`). Embed the markdown directly in your reply to display the image to the user.",
    parameters: {
      imageURLs: z
        .array(httpUrlSchema)
        .optional()
        .describe("HTTP(S) URLs of images to download. Local file paths are rejected."),
      websiteURL: httpUrlSchema.optional().describe("URL of a page to scrape <img> tags from and download."),
    },

    /**
     * Execute an image download batch, optionally preceded by scraping image URLs from a page.
     *
     * @param arguments_ - Validated tool parameters.
     * @param context - Runtime tool context supplied by the SDK.
     * @returns Per-image records carrying filename, alt, title, and either a markdown reference or an error, or a user-facing error string.
     */
    implementation: async (arguments_, context) => {
      const { imageURLs, websiteURL } = arguments_
      const explicitUrls = imageURLs ?? []
      const shouldScrapePage = websiteURL !== undefined

      if (explicitUrls.length === 0 && !shouldScrapePage) {
        return "Error: Provide at least one of imageURLs or websiteURL."
      }

      try {
        const { maxImages, maxResponseBytes, maxImageBytes } = resolveConfig(ctl)
        const subjects: ImageSubject[] = explicitUrls.map(source => ({ src: source, alt: "", title: "" }))

        if (shouldScrapePage) {
          context.status("Fetching image URLs from website...")

          const page = await fetchPage(impit, pageCache, websiteURL, {
            signal: context.signal,
            retry,
            onFailedAttempt: createRetryNotifier(context.status, "website fetch"),
            maxBytes: maxResponseBytes,
            limiter: rateLimiter,
          })

          if (page.kind !== "html") {
            throw new UnsupportedContentTypeError(page.mimeType, websiteURL)
          }

          const scraped = extractPageImages(parseHtml(page.html), websiteURL, maxImages)

          for (const image of scraped) {
            subjects.push({ src: image.src, alt: image.alt, title: image.title })
          }
        }

        if (subjects.length === 0) {
          context.warn("No supported images found on page")

          return []
        }

        context.status("Downloading images...")

        const batch = await downloadImages(
          subjects.map(subject => subject.src),
          impit,
          { workingDirectory: ctl.getWorkingDirectory(), timestamp: Date.now(), maxBytes: maxImageBytes },
          {
            warn: context.warn,
            signal: context.signal,
            limiter: imageLimiter,
            retry,
            onFailedAttempt: createRetryNotifier(context.status, "image download"),
          }
        )
        const rendered = renderImageResults(subjects, batch)

        context.status(`Processed ${rendered.length} images.`)

        return rendered
      } catch (error) {
        return formatToolError(error, context, "image-download")
      }
    },
  })
}