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

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_MAX_ATTEMPTS = 3
QUERIT_RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504}
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:
        headers = {
            "Accept": "application/json",
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        for attempt in range(QUERIT_MAX_ATTEMPTS):
            if self.check_if_canceled("QueritSearch processing"):
                raise _QueritCanceled
            try:
                response = requests.post(
                    QUERIT_SEARCH_URL,
                    headers=headers,
                    json=payload,
                    timeout=DEFAULT_TIMEOUT,
                )
                if response.status_code in QUERIT_RETRYABLE_STATUS_CODES and attempt + 1 < QUERIT_MAX_ATTEMPTS:
                    self._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
                self._wait_before_retry()
        raise RuntimeError("Querit request failed after three attempts.")

    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", "-_-!"))


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 _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)