爬取网页数据

package nexus.io.mosskb.service.kb;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.function.Consumer;

import org.jsoup.Connection;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;

public class WebCrawlService {

  /** 部分站点会过滤默认 UA，因此使用浏览器 UA。 */
  public static final String USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36";

  private static final int TIMEOUT_MS = 20_000;
  private static final int MAX_BODY_BYTES = 5 * 1024 * 1024;
  private static final int MAX_DEPTH = 2;
  private static final int MAX_PAGES = 50;

  public record Page(String url, String name, String markdown, List<String> links) {
  }

  /** 抓取单个页面，链接只用于继续扩散，不计入正文。 */
  public Page fetch(String url, String selector) throws IOException {
    String target = absoluteUrl(url);
    if (target == null) {
      throw new CrawlException("Web 根地址必须是 http 或 https 地址");
    }
    Connection.Response response;
    try {
      response = Jsoup.connect(target).userAgent(USER_AGENT).timeout(TIMEOUT_MS).maxBodySize(MAX_BODY_BYTES)
          .followRedirects(true).ignoreHttpErrors(true).execute();
    } catch (IOException e) {
      throw new CrawlException("无法访问该地址：" + e.getMessage());
    }
    if (response.statusCode() != 200) {
      throw new CrawlException("目标地址返回 HTTP " + response.statusCode());
    }
    Document document = response.parse();
    // 记录重定向之后的正式地址，避免同一页面按多个地址重复入库。
    String finalUrl = absoluteUrl(response.url() == null ? target : response.url().toString());
    return new Page(finalUrl == null ? target : finalUrl, pageName(document, finalUrl), markdown(document, selector),
        pageLinks(document));
  }
}

package nexus.io.mosskb.service.kb;

import java.io.IOException;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.LinkedHashSet;
import java.util.Set;
import java.util.function.Consumer;

public class WebCrawlService {

  /**
   * @param firstPage 已经抓取过的入口页面，可为空；不为空时不再重复请求入口地址
   * @param onPage    每抓到一个页面就回调一次，由调用方决定如何落库
   */
  public int crawl(String sourceUrl, String selector, Page firstPage, Consumer<Page> onPage) throws IOException {
    String entry = absoluteUrl(sourceUrl);
    if (entry == null) {
      throw new CrawlException("Web 根地址必须是 http 或 https 地址");
    }
    String prefix = directoryPrefix(entry);
    Set<String> visited = new LinkedHashSet<>();
    Deque<Object[]> queue = new ArrayDeque<>();
    int count = 0;
    if (firstPage != null) {
      count++;
      visited.add(trimSlash(firstPage.url()));
      onPage.accept(firstPage);
      for (String link : firstPage.links()) {
        if (link.startsWith(prefix)) {
          queue.add(new Object[] { link, 1 });
        }
      }
    } else {
      queue.add(new Object[] { entry, 0 });
    }
    while (!queue.isEmpty() && count < MAX_PAGES) {
      Object[] current = queue.poll();
      String url = (String) current[0];
      int depth = (Integer) current[1];
      if (url == null || !visited.add(trimSlash(url))) {
        continue;
      }
      Page page;
      try {
        page = fetch(url, selector);
      } catch (IOException e) {
        continue;
      }
      String canonicalKey = trimSlash(page.url());
      if (!canonicalKey.equals(trimSlash(url)) && !visited.add(canonicalKey)) {
        continue;
      }
      count++;
      onPage.accept(page);
      if (depth >= MAX_DEPTH) {
        continue;
      }
      for (String link : page.links()) {
        if (link.startsWith(prefix) && !visited.contains(trimSlash(link))) {
          queue.add(new Object[] { link, depth + 1 });
        }
      }
    }
    return count;
  }
}

package nexus.io.mosskb.service.kb;

import java.util.Locale;
import java.util.Set;

import org.jsoup.nodes.Element;
import org.jsoup.nodes.Node;
import org.jsoup.nodes.TextNode;

public class WebCrawlService {

  private static final Set<String> SKIPPED_TAGS = Set.of("script", "style", "noscript", "template", "svg", "canvas",
      "iframe", "head", "form", "select", "option", "button", "nav", "footer", "aside");

  private static void appendNode(Node node, StringBuilder out, String baseUri) {
    if (node instanceof TextNode textNode) {
      out.append(textNode.getWholeText().replace('\u00a0', ' '));
      return;
    }
    if (!(node instanceof Element element)) {
      return;
    }
    String tag = element.tagName().toLowerCase(Locale.ROOT);
    if (SKIPPED_TAGS.contains(tag) || element.hasAttr("hidden")) {
      return;
    }
    switch (tag) {
      case "h1", "h2", "h3", "h4", "h5", "h6" -> {
        int level = tag.charAt(1) - '0';
        out.append("\n\n").append("#".repeat(level)).append(' ').append(inline(element, baseUri)).append("\n\n");
      }
      case "p", "div", "section", "article", "main", "header" -> {
        out.append("\n\n");
        appendChildren(element, out, baseUri);
        out.append("\n\n");
      }
      case "ul", "ol" -> appendList(element, out, baseUri, tag);
      case "pre" -> out.append("\n\n```\n").append(element.wholeText().strip()).append("\n```\n\n");
      case "table" -> out.append("\n\n").append(table(element, baseUri)).append("\n\n");
      case "a" -> appendLink(element, out, baseUri);
      case "img" -> appendImage(element, out);
      default -> appendChildren(element, out, baseUri);
    }
  }
}

package nexus.io.mosskb.service.kb;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class WebCrawlService {

  private static final Pattern HEADING = Pattern.compile("^(#{1,6})\\s+(.*)$");

  public record Paragraph(String title, String content) {
  }

  /** 按标题切分 Markdown：标题构成标题链，正文构成分段内容。 */
  public List<Paragraph> split(String markdown) {
    List<Paragraph> paragraphs = new ArrayList<>();
    if (markdown == null || markdown.isBlank()) {
      return paragraphs;
    }
    String normalized = markdown.replace("\r\n", "\n").replace('\r', '\n');
    List<String> chain = new ArrayList<>();
    List<String> sectionTitles = List.of();
    StringBuilder body = new StringBuilder();
    for (String line : normalized.split("\n", -1)) {
      Matcher matcher = HEADING.matcher(line.strip());
      if (!matcher.matches()) {
        body.append(line).append('\n');
        continue;
      }
      appendSection(paragraphs, sectionTitles, body);
      body.setLength(0);
      int level = matcher.group(1).length();
      while (chain.size() >= level) {
        chain.remove(chain.size() - 1);
      }
      chain.add(matcher.group(2).strip());
      sectionTitles = List.copyOf(chain);
    }
    appendSection(paragraphs, sectionTitles, body);
    return paragraphs;
  }
}

package nexus.io.mosskb.service.kb;

public class MossKbWebDatasetService {

  /** 同一个地址只对应一个文档，替换同步据此复用已有文档。 */
  private static final String FIND_DOCUMENT_BY_SOURCE_URL = """
      select id
        from moss_kb_document
       where dataset_id = ?
         and meta->>'source_url' = ?
      """;

  private static final String DELETE_DOCUMENT_PARAGRAPHS = """
      delete from moss_kb_paragraph
       where document_id = ?
      """;
}

package nexus.io.mosskb.service.kb;

import java.util.ArrayList;
import java.util.List;

import org.postgresql.util.PGobject;

import com.jfinal.kit.Kv;

import nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.constant.MossKbTableNames;
import nexus.io.tio.utils.crypto.Md5Utils;
import nexus.io.tio.utils.snowflake.SnowflakeIdUtils;

public class MossKbWebDatasetService {

  public static final String DOCUMENT_TYPE_WEB = "1";

  /** 写入或替换一个网页文档，段落与文档在同一个事务里落库。 */
  private int saveDocument(WebCrawlService crawler, Long datasetId, Long userId, Long embeddingModeId, String selector,
      WebCrawlService.Page page, Long documentId) {
    List<WebCrawlService.Paragraph> paragraphs = crawler.split(page.markdown());
    if (paragraphs.isEmpty()) {
      throw new IllegalArgumentException("页面没有可用的正文");
    }
    long id = documentId == null ? SnowflakeIdUtils.id() : documentId;
    KbEmbeddingService embeddingService = Aop.get(KbEmbeddingService.class);
    List<Row> rows = new ArrayList<>();
    int charLength = 0;
    for (WebCrawlService.Paragraph paragraph : paragraphs) {
      charLength += paragraph.content().length();
      PGobject titleVector = paragraph.title().isBlank() ? null
          : embeddingService.getVectorForModel(paragraph.title(), embeddingModeId);
      rows.add(Row.by("id", SnowflakeIdUtils.id()).set("source_id", id).set("source_type", DOCUMENT_TYPE_WEB)
          .set("title", paragraph.title()).set("content", paragraph.content())
          .set("md5", Md5Utils.md5Hex(paragraph.content())).set("status", "2").set("hit_num", 0).set("is_active", true)
          .set("dataset_id", datasetId).set("document_id", id).set("embedding",
              embeddingService.getVectorForModel(paragraph.content(), embeddingModeId))
          .set("title_embedding", titleVector));
    }
    Row document = Row.by("id", id).set("user_id", userId).set("name", page.name()).set("type", DOCUMENT_TYPE_WEB)
        .set("char_length", charLength).set("status", "2").set("is_active", true).set("dataset_id", datasetId)
        .set("paragraph_count", rows.size()).set("hit_handling_method", "optimization")
        .set("directly_return_similarity", 0.9).set("meta", pageMeta(page.url(), selector));
    Db.tx(() -> {
      if (documentId == null) {
        Db.save(MossKbTableNames.moss_kb_document, document);
      } else {
        document.set("update_time", new java.util.Date());
        Db.update(MossKbTableNames.moss_kb_document, document);
        Db.update(DELETE_DOCUMENT_PARAGRAPHS, id);
      }
      Db.batchSave(MossKbTableNames.moss_kb_paragraph, rows, 2000);
      return true;
    });
    return rows.size();
  }
}

package nexus.io.mosskb.service.kb;

import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import lombok.extern.slf4j.Slf4j;
import nexus.io.mosskb.utils.ExecutorServiceUtils;

@Slf4j
public class MossKbWebDatasetService {

  private static final Set<String> RUNNING = ConcurrentHashMap.newKeySet();
  private static final String DATASET_KEY = "dataset:";
  private static final String DOCUMENT_KEY = "document:";

  /** 占用任务槽；返回 false 说明同一个知识库或文档的任务还在执行。 */
  private boolean acquire(String key) {
    return RUNNING.add(key);
  }

  private void release(String key) {
    RUNNING.remove(key);
  }

  /** 提交任务，任务结束（含异常）后释放任务槽。 */
  private void submit(String key, Runnable task) {
    submit(() -> {
      try {
        task.run();
      } finally {
        release(key);
      }
    });
  }

  private void submit(Runnable task) {
    ExecutorServiceUtils.getExecutorService().submit(() -> {
      try {
        task.run();
      } catch (Exception e) {
        log.error("Web 站点任务执行失败", e);
      }
    });
  }
}

package nexus.io.mosskb.controller;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;

import nexus.io.annotation.Post;
import nexus.io.annotation.Put;
import nexus.io.annotation.RequestPath;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.service.kb.MossKbWebDatasetService;
import nexus.io.model.result.ResultVo;
import nexus.io.tio.boot.http.TioRequestContext;
import nexus.io.tio.http.common.HttpRequest;

@RequestPath("/api/dataset")
public class ApiDatasetController {

  /** 创建 Web 站点知识库：保存根地址与选择器后由后台线程抓取页面。 */
  @Post("/web")
  public ResultVo saveWebDataset(HttpRequest request) {
    JSONObject input = JSON.parseObject(request.getBodyString());
    return Aop.get(MossKbWebDatasetService.class).save(TioRequestContext.getUserIdLong(), input);
  }

  /** 同步 Web 站点知识库，sync_type 取 replace 或 complete。 */
  @Put("/{id}/sync_web")
  public ResultVo syncWebDataset(Long id, HttpRequest request) {
    return Aop.get(MossKbWebDatasetService.class).sync(TioRequestContext.getUserIdLong(), id,
        request.getParam("sync_type"));
  }

  /** 按地址列表向 Web 站点知识库导入网页文档。 */
  @Post("/{datasetId}/document/web")
  public ResultVo importWebDocument(Long datasetId, HttpRequest request) {
    JSONObject input = JSON.parseObject(request.getBodyString());
    return Aop.get(MossKbWebDatasetService.class).importDocuments(TioRequestContext.getUserIdLong(), datasetId, input);
  }
}