ragflow/agent/tools/querit.py

#
#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import logging
import os
import re
import time
from abc import ABC
from typing import Any
from urllib.parse import urlparse

import requests

from agent.tools.base import ToolBase, ToolMeta, ToolParamBase
from common.connection_utils import timeout
from common.http_client import DEFAULT_TIMEOUT

QUERIT_SEARCH_URL = "https://api.querit.ai/v1/search"
QUERIT_CONTENTS_URL = "https://api.querit.ai/v1/contents"
QUERIT_MAX_ATTEMPTS = 3
QUERIT_RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504}
QUERIT_CONTENT_FORMATS = {"text", "markdown", "html"}
TIME_RANGE_PATTERN = re.compile(r"^([dwmy][1-9][0-9]*|\d{4}-\d{2}-\d{2}to\d{4}-\d{2}-\d{2})$")
logger = logging.getLogger(__name__)


class _QueritCanceled(Exception):
    pass


class QueritSearchParam(ToolParamBase):
    def __init__(self):
        self.meta: ToolMeta = {
            "name": "querit_search",
            "description": "Search the live web with Querit and return the complete Querit API response.",
            "parameters": {
                "query": {
                    "type": "string",
                    "description": "The search query to execute with Querit.",
                    "default": "{sys.query}",
                    "required": True,
                },
                "count": {
                    "type": "integer",
                    "description": "The maximum number of results to return. Defaults to 10.",
                    "default": 10,
                    "required": False,
                },
                "chunks_per_doc": {
                    "type": "integer",
                    "description": "The number of summary chunks per document. Supports values from 1 to 3.",
                    "default": 3,
                    "required": False,
                },
                "site_include": {
                    "type": "array",
                    "description": "Sites that search results must include.",
                    "default": [],
                    "items": {"type": "string"},
                    "required": False,
                },
                "site_exclude": {
                    "type": "array",
                    "description": "Sites that search results must exclude.",
                    "default": [],
                    "items": {"type": "string"},
                    "required": False,
                },
                "time_range": {
                    "type": "string",
                    "description": "A Querit time range such as d7, w1, m3, y1, or YYYY-MM-DDtoYYYY-MM-DD.",
                    "default": "",
                    "required": False,
                },
                "country_include": {
                    "type": "array",
                    "description": "Return results associated with the specified countries.",
                    "default": [],
                    "items": {"type": "string"},
                    "required": False,
                },
                "language_include": {
                    "type": "array",
                    "description": "Languages that search results must include.",
                    "default": [],
                    "items": {"type": "string"},
                    "required": False,
                },
            },
        }
        super().__init__()
        self.api_key = ""

    def check(self):
        _validate_search_inputs(
            count=self.count,
            chunks_per_doc=self.chunks_per_doc,
            time_range=self.time_range,
            site_include=self.site_include,
            site_exclude=self.site_exclude,
            country_include=self.country_include,
            language_include=self.language_include,
        )

    def get_input_form(self) -> dict[str, dict]:
        return {
            "query": {"name": "Query", "type": "line"},
            "count": {"name": "Count", "type": "line"},
            "chunks_per_doc": {"name": "Chunks per document", "type": "line"},
            "site_include": {"name": "Include sites", "type": "line"},
            "site_exclude": {"name": "Exclude sites", "type": "line"},
            "time_range": {"name": "Time range", "type": "line"},
            "country_include": {"name": "Include countries", "type": "line"},
            "language_include": {"name": "Include languages", "type": "line"},
        }


class QueritSearch(ToolBase, ABC):
    component_name = "QueritSearch"

    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", "12")))
    def _invoke(self, **kwargs):
        if self.check_if_canceled("QueritSearch processing"):
            return

        query = kwargs.get("query")
        if not isinstance(query, str):
            return self._fail("Querit query must be a string.")
        if not query:
            self.set_output("formalized_content", "")
            self.set_output("json", {})
            return ""

        node_api_key = (self._param.api_key or "").strip()
        api_key = node_api_key or (os.environ.get("QUERIT_API_KEY") or "").strip()
        if not api_key:
            return self._fail("Querit API key is required. Configure api_key or set QUERIT_API_KEY.")

        values = {
            name: kwargs[name] if name in kwargs else getattr(self._param, name)
            for name in (
                "count",
                "chunks_per_doc",
                "site_include",
                "site_exclude",
                "time_range",
                "country_include",
                "language_include",
            )
        }

        try:
            _validate_search_inputs(**values)
            payload = _build_payload(query, **values)
            response_data = self._search(payload, api_key)
            if not isinstance(response_data, dict):
                raise TypeError("Querit API response must be a JSON object.")

            result_container = response_data.get("results", {})
            if not isinstance(result_container, dict):
                raise TypeError("Querit API response field results must be an object.")
            results = result_container.get("result", [])
            if not isinstance(results, list):
                raise TypeError("Querit API response field results.result must be an array.")

            reference_results = [item for item in results if isinstance(item, dict)]
            if reference_results:
                self._retrieve_chunks(
                    reference_results,
                    get_title=lambda item: _querit_text(item.get("title")),
                    get_url=lambda item: _querit_text(item.get("url")),
                    get_content=lambda item: _querit_text(item.get("snippet")),
                    get_score=lambda _item: 1,
                )
            else:
                self.set_output("formalized_content", "")
            self.set_output("json", response_data)
            return self.output("formalized_content")
        except _QueritCanceled:
            return
        except (requests.RequestException, RuntimeError, TypeError, ValueError) as error:
            return self._fail(_safe_error_message(error, api_key))

    def _search(self, payload: dict[str, Any], api_key: str) -> Any:
        return _post_querit(self, QUERIT_SEARCH_URL, payload, api_key, "QueritSearch")

    def _wait_before_retry(self) -> None:
        if self.check_if_canceled("QueritSearch processing"):
            raise _QueritCanceled
        time.sleep(self._param.delay_after_error)

    def _fail(self, message: str) -> str:
        self.set_output("_ERROR", message)
        logger.error("Querit search failed: %s", message)
        return f"Querit error: {message}"

    def thoughts(self) -> str:
        return "Searching Querit for `{}`.".format(self.get_input().get("query", "-_-!"))


class QueritContentsParam(ToolParamBase):
    def __init__(self):
        self.meta: ToolMeta = {
            "name": "querit_contents",
            "description": "Crawl one or more web pages with Querit and return their contents.",
            "parameters": {
                "urls": {
                    "type": "array",
                    "description": "The absolute HTTP or HTTPS URLs to crawl. Supports 1 to 10 URLs.",
                    "default": [],
                    "items": {"type": "string"},
                    "required": True,
                },
                "format": {
                    "type": "string",
                    "description": "Content format: text, markdown, or html. Defaults to markdown.",
                    "enum": ["text", "markdown", "html"],
                    "default": "markdown",
                    "required": False,
                },
                "crawl_timeout": {
                    "type": "integer",
                    "description": "Per-page crawl timeout in seconds. Must be between 1 and 60.",
                    "default": 10,
                    "required": False,
                },
                "extras_meta": {
                    "type": "boolean",
                    "description": "Whether to include page metadata in each result.",
                    "default": False,
                    "required": False,
                },
            },
        }
        super().__init__()
        self.api_key = ""

    def check(self):
        self.urls = _normalize_contents_urls(self.urls)
        _validate_contents_inputs(self.urls, self.format, self.crawl_timeout, self.extras_meta)

    def get_input_form(self) -> dict[str, dict]:
        return {"urls": {"name": "URLs", "type": "line"}}


class QueritContents(ToolBase, ABC):
    component_name = "QueritContents"

    @timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", "70")))
    def _invoke(self, **kwargs):
        if self.check_if_canceled("QueritContents processing"):
            return

        values = {name: kwargs[name] if name in kwargs else getattr(self._param, name) for name in ("urls", "format", "crawl_timeout", "extras_meta")}
        values["urls"] = _normalize_contents_urls(values["urls"])

        node_api_key = (self._param.api_key or "").strip()
        api_key = node_api_key or (os.environ.get("QUERIT_API_KEY") or "").strip()
        if not api_key:
            return self._fail("Querit API key is required. Configure api_key or set QUERIT_API_KEY.")

        try:
            _validate_contents_inputs(**values)
            response_data = self._request(_build_contents_payload(**values), api_key)
            _validate_contents_response(response_data)
            self.set_output("json", response_data)
            return self.output("json")
        except _QueritCanceled:
            return
        except (requests.RequestException, RuntimeError, TypeError, ValueError) as error:
            return self._fail(_safe_error_message(error, api_key))

    def _request(self, payload: dict[str, Any], api_key: str) -> Any:
        request_timeout = max(DEFAULT_TIMEOUT, payload["crawlTimeout"] + 5)
        return _post_querit(self, QUERIT_CONTENTS_URL, payload, api_key, "QueritContents", request_timeout)

    def _wait_before_retry(self) -> None:
        if self.check_if_canceled("QueritContents processing"):
            raise _QueritCanceled
        time.sleep(self._param.delay_after_error)

    def _fail(self, message: str) -> str:
        self.set_output("_ERROR", message)
        logger.error("Querit contents failed: %s", message)
        return f"Querit contents error: {message}"

    def thoughts(self) -> str:
        return "Reading web page contents with Querit."


def _build_payload(query: str, **values: Any) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "query": query,
        "count": values["count"],
    }
    if values["chunks_per_doc"] is not None:
        payload["chunksPerDoc"] = values["chunks_per_doc"]

    filters: dict[str, Any] = {}
    if values["site_include"] or values["site_exclude"]:
        filters["sites"] = {}
        if values["site_include"]:
            filters["sites"]["include"] = values["site_include"]
        if values["site_exclude"]:
            filters["sites"]["exclude"] = values["site_exclude"]
    if values["time_range"]:
        filters["timeRange"] = {"date": values["time_range"]}
    if values["country_include"]:
        filters["geo"] = {"countries": {"include": values["country_include"]}}
    if values["language_include"]:
        filters["languages"] = {"include": values["language_include"]}
    if filters:
        payload["filters"] = filters
    return payload


def _post_querit(tool: Any, endpoint: str, payload: dict[str, Any], api_key: str, operation: str, request_timeout: float = DEFAULT_TIMEOUT) -> Any:
    headers = {
        "Accept": "application/json",
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    for attempt in range(QUERIT_MAX_ATTEMPTS):
        if tool.check_if_canceled(f"{operation} processing"):
            raise _QueritCanceled
        try:
            response = requests.post(endpoint, headers=headers, json=payload, timeout=request_timeout)
            if response.status_code in QUERIT_RETRYABLE_STATUS_CODES and attempt + 1 < QUERIT_MAX_ATTEMPTS:
                tool._wait_before_retry()
                continue
            response.raise_for_status()
            return response.json()
        except requests.JSONDecodeError:
            raise
        except requests.HTTPError:
            raise
        except requests.RequestException:
            if attempt + 1 >= QUERIT_MAX_ATTEMPTS:
                raise
            tool._wait_before_retry()
    raise RuntimeError("Querit request failed after three attempts.")


def _build_contents_payload(urls: list[str], format: str, crawl_timeout: int, extras_meta: bool) -> dict[str, Any]:
    return {
        "urls": urls,
        "format": format,
        "crawlTimeout": crawl_timeout,
        "extrasMeta": extras_meta,
    }


def _normalize_contents_urls(urls: Any) -> Any:
    if isinstance(urls, str):
        return [url.strip() for url in urls.split(",") if url.strip()]
    return urls


def _validate_contents_inputs(urls: Any, format: Any, crawl_timeout: Any, extras_meta: Any) -> None:
    if not isinstance(urls, list) or not 1 <= len(urls) <= 10 or any(not isinstance(url, str) or not url.strip() for url in urls):
        raise ValueError("Querit urls must contain between 1 and 10 non-empty strings.")
    for url in urls:
        parsed = urlparse(url)
        if parsed.scheme not in {"http", "https"} or not parsed.netloc:
            raise ValueError("Querit urls must be absolute HTTP or HTTPS URLs.")
    if format not in QUERIT_CONTENT_FORMATS:
        raise ValueError("Querit format must be text, markdown, or html.")
    if type(crawl_timeout) is not int or not 1 <= crawl_timeout <= 60:
        raise ValueError("Querit crawl_timeout must be an integer from 1 to 60.")
    if type(extras_meta) is not bool:
        raise ValueError("Querit extras_meta must be a boolean.")


def _validate_contents_response(response_data: Any) -> None:
    if not isinstance(response_data, dict):
        raise TypeError("Querit API response must be a JSON object.")
    if "results" in response_data and not isinstance(response_data["results"], list):
        raise TypeError("Querit API response field results must be an array.")
    if "statuses" in response_data and not isinstance(response_data["statuses"], list):
        raise TypeError("Querit API response field statuses must be an array.")


def _validate_search_inputs(
    count: Any,
    chunks_per_doc: Any,
    time_range: Any,
    site_include: Any,
    site_exclude: Any,
    country_include: Any,
    language_include: Any,
) -> None:
    if type(count) is not int or count < 1:
        raise ValueError("Querit count must be an integer greater than or equal to 1.")
    if chunks_per_doc is not None and (type(chunks_per_doc) is not int or not 1 <= chunks_per_doc <= 3):
        raise ValueError("Querit chunks_per_doc must be an integer from 1 to 3.")
    if type(time_range) is not str:
        raise ValueError("Querit time_range must be a string.")
    if time_range and not TIME_RANGE_PATTERN.fullmatch(time_range):
        raise ValueError("Querit time_range must use dN, wN, mN, yN, or YYYY-MM-DDtoYYYY-MM-DD.")
    for name, value in (
        ("site_include", site_include),
        ("site_exclude", site_exclude),
        ("country_include", country_include),
        ("language_include", language_include),
    ):
        if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
            raise ValueError(f"Querit {name} must be an array of strings.")


def _safe_error_message(error: Exception, api_key: str) -> str:
    message = str(error) or error.__class__.__name__
    return message.replace(api_key, "[REDACTED]") if api_key else message


def _querit_text(value: Any) -> str:
    return "" if value is None else str(value)