mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-30 12:39:27 +08:00
feat(agent): add Querit search tool (#17548)
This commit is contained in:
298
agent/tools/querit.py
Normal file
298
agent/tools/querit.py
Normal file
@@ -0,0 +1,298 @@
|
||||
#
|
||||
# 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)
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
@@ -69,7 +70,12 @@ func (c *ToolBackedComponent) Invoke(ctx context.Context, db *gorm.DB, inputs ma
|
||||
}
|
||||
|
||||
raw, invokeErr := c.tool.InvokableRun(ctx, string(argsJSON))
|
||||
decoded := parseToolEnvelope(raw)
|
||||
var decoded map[string]any
|
||||
if c.spec.PreserveJSONNumbers {
|
||||
decoded = parseToolEnvelopeLossless(raw)
|
||||
} else {
|
||||
decoded = parseToolEnvelope(raw)
|
||||
}
|
||||
if rawValue, invalid := decoded["_raw"]; invalid {
|
||||
if invokeErr != nil {
|
||||
return nil, fmt.Errorf("canvas: %s: %w", c.name, invokeErr)
|
||||
@@ -97,6 +103,19 @@ func (c *ToolBackedComponent) Invoke(ctx context.Context, db *gorm.DB, inputs ma
|
||||
return c.tool.BuildComponentOutputs(decoded), nil
|
||||
}
|
||||
|
||||
func parseToolEnvelopeLossless(jsonStr string) map[string]any {
|
||||
var out map[string]any
|
||||
decoder := json.NewDecoder(strings.NewReader(jsonStr))
|
||||
decoder.UseNumber()
|
||||
if err := decoder.Decode(&out); err != nil {
|
||||
return map[string]any{"_raw": jsonStr}
|
||||
}
|
||||
if err := decoder.Decode(&struct{}{}); err != io.EOF {
|
||||
return map[string]any{"_raw": jsonStr}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (c *ToolBackedComponent) Stream(_ context.Context, _ *gorm.DB, _ map[string]any) (<-chan map[string]any, error) {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -115,6 +134,7 @@ var toolComponentRegistrations = []struct {
|
||||
{componentName: "GoogleScholar", toolName: "google_scholar"},
|
||||
{componentName: "KeenableSearch", toolName: "keenable"},
|
||||
{componentName: "PubMed", toolName: "pubmed"},
|
||||
{componentName: "QueritSearch", toolName: "querit_search"},
|
||||
{componentName: "SearXNG", toolName: "searxng"},
|
||||
{componentName: "TavilySearch", toolName: "tavily"},
|
||||
{componentName: "TavilyExtract", toolName: "tavily_extract"},
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"net/url"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
einotool "github.com/cloudwego/eino/components/tool"
|
||||
@@ -192,6 +193,13 @@ func TestToolBackedComponentRegisteredFactories(t *testing.T) {
|
||||
outputKey: "success",
|
||||
inputKey: "to_email",
|
||||
},
|
||||
{
|
||||
name: "QueritSearch",
|
||||
toolName: "QueritSearch",
|
||||
params: map[string]any{"api_key": "stored-key", "count": float64(10), "chunks_per_doc": float64(3), "outputs": map[string]any{"json": map[string]any{}}},
|
||||
outputKey: "json",
|
||||
inputKey: "query",
|
||||
},
|
||||
{
|
||||
name: "SearXNG",
|
||||
toolName: "SearXNG",
|
||||
@@ -275,7 +283,7 @@ func TestToolBackedComponentWenCaiInvoke(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestToolBackedComponentRegisteredBuildWorkflow(t *testing.T) {
|
||||
for _, componentName := range []string{"ArXiv", "BGPT", "DuckDuckGo", "Email", "Google", "GoogleScholar", "KeenableSearch", "PubMed", "SearXNG", "WenCai", "TavilyExtract", "TavilySearch", "Wikipedia", "YahooFinance"} {
|
||||
for _, componentName := range []string{"ArXiv", "BGPT", "DuckDuckGo", "Email", "Google", "GoogleScholar", "KeenableSearch", "PubMed", "QueritSearch", "SearXNG", "WenCai", "TavilyExtract", "TavilySearch", "Wikipedia", "YahooFinance"} {
|
||||
t.Run(componentName, func(t *testing.T) {
|
||||
c := &canvas.Canvas{
|
||||
Components: map[string]canvas.CanvasComponent{
|
||||
@@ -468,6 +476,65 @@ func TestToolBackedComponentTavilyIntegration(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolBackedComponentQueritIntegration(t *testing.T) {
|
||||
var serverCalls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
serverCalls.Add(1)
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[{"title":"RAGFlow","url":"https://ragflow.io","snippet":"RAG article","custom":"preserved"}]},"search_id":11099848653006015581,"query_context":{"rewritten":"rag flow"}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
target, err := url.Parse(server.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("parse test server URL: %v", err)
|
||||
}
|
||||
helper := agenttool.NewHTTPHelper().WithClient(&http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) {
|
||||
cloned := request.Clone(request.Context())
|
||||
cloned.URL.Scheme = target.Scheme
|
||||
cloned.URL.Host = target.Host
|
||||
return http.DefaultTransport.RoundTrip(cloned)
|
||||
})})
|
||||
querit := agenttool.NewQueritToolWithEnvKey(helper, func() string { return "key" })
|
||||
component := &ToolBackedComponent{name: "QueritSearch", tool: querit, spec: querit.ComponentSpec()}
|
||||
state := runtime.NewCanvasState("run-querit", "task-querit")
|
||||
empty, err := component.Invoke(context.Background(), nil, map[string]any{"query": ""})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke(empty query): %v", err)
|
||||
}
|
||||
emptyJSON, emptyJSONOK := empty["json"].(map[string]any)
|
||||
if serverCalls.Load() != 0 || empty["formalized_content"] != "" || !emptyJSONOK || len(emptyJSON) != 0 {
|
||||
t.Fatalf("empty query result = %#v, server calls = %d", empty, serverCalls.Load())
|
||||
}
|
||||
out, err := component.Invoke(runtime.WithState(context.Background(), state), nil, map[string]any{"query": "ragflow"})
|
||||
if err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
rendered, _ := out["formalized_content"].(string)
|
||||
if !strings.Contains(rendered, "Title: RAGFlow") || !strings.Contains(rendered, "RAG article") {
|
||||
t.Fatalf("formalized_content = %q", rendered)
|
||||
}
|
||||
complete, ok := out["json"].(map[string]any)
|
||||
searchID, searchIDOK := complete["search_id"].(json.Number)
|
||||
if !ok || !searchIDOK || searchID.String() != "11099848653006015581" || complete["query_context"] == nil {
|
||||
t.Fatalf("complete json = %#v", out["json"])
|
||||
}
|
||||
encoded, err := json.Marshal(complete)
|
||||
if err != nil || !strings.Contains(string(encoded), `"search_id":11099848653006015581`) {
|
||||
t.Fatalf("re-encoded complete json = %s, %v", encoded, err)
|
||||
}
|
||||
chunks := state.GetRetrievalChunks()
|
||||
if len(chunks) != 1 || chunks[0]["document_name"] != "RAGFlow" || chunks[0]["similarity"] != 1 {
|
||||
t.Fatalf("recorded references = %#v", chunks)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseToolEnvelopeLosslessRejectsTrailingContent(t *testing.T) {
|
||||
out := parseToolEnvelopeLossless(`{"results":{"result":[]}} trailing`)
|
||||
if out["_raw"] == nil {
|
||||
t.Fatalf("trailing content was accepted: %#v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolBackedComponentYahooFinanceIntegration(t *testing.T) {
|
||||
serverCalls := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
|
||||
605
internal/agent/tool/querit.go
Normal file
605
internal/agent/tool/querit.go
Normal file
@@ -0,0 +1,605 @@
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha1"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
|
||||
const (
|
||||
queritToolName = "querit_search"
|
||||
queritToolDescription = "Search the web via the Querit API and return the complete search response."
|
||||
queritEndpoint = "https://api.querit.ai/v1/search"
|
||||
queritMaxAttempts = 3
|
||||
queritPromptMaxTokens = 200000
|
||||
)
|
||||
|
||||
var (
|
||||
queritRelativeTimeRangePattern = regexp.MustCompile(`^[dwmy][1-9][0-9]*$`)
|
||||
queritDateRangePattern = regexp.MustCompile(`^\d{4}-\d{2}-\d{2}to\d{4}-\d{2}-\d{2}$`)
|
||||
queritDataImagePattern = regexp.MustCompile(`!?\[[a-z]+\]\(data:image/png;base64,[ 0-9A-Za-z/_=+\-]+\)`)
|
||||
queritNewlinePattern = regexp.MustCompile(`\n+`)
|
||||
)
|
||||
|
||||
type queritParams struct {
|
||||
APIKey string `json:"api_key"`
|
||||
Query string `json:"query"`
|
||||
Count int `json:"count"`
|
||||
ChunksPerDoc *int `json:"chunks_per_doc"`
|
||||
SiteInclude []string `json:"site_include"`
|
||||
SiteExclude []string `json:"site_exclude"`
|
||||
TimeRange string `json:"time_range"`
|
||||
CountryInclude []string `json:"country_include"`
|
||||
LanguageInclude []string `json:"language_include"`
|
||||
}
|
||||
|
||||
type queritRequest struct {
|
||||
Query string `json:"query"`
|
||||
Count int `json:"count"`
|
||||
ChunksPerDoc *int `json:"chunksPerDoc,omitempty"`
|
||||
Filters *queritFilters `json:"filters,omitempty"`
|
||||
}
|
||||
|
||||
type queritFilters struct {
|
||||
Sites *queritIncludeExcludeFilter `json:"sites,omitempty"`
|
||||
TimeRange *queritDateFilter `json:"timeRange,omitempty"`
|
||||
Geo *queritGeoFilter `json:"geo,omitempty"`
|
||||
Languages *queritIncludeFilter `json:"languages,omitempty"`
|
||||
}
|
||||
|
||||
type queritIncludeExcludeFilter struct {
|
||||
Include []string `json:"include,omitempty"`
|
||||
Exclude []string `json:"exclude,omitempty"`
|
||||
}
|
||||
|
||||
type queritIncludeFilter struct {
|
||||
Include []string `json:"include,omitempty"`
|
||||
}
|
||||
|
||||
type queritDateFilter struct {
|
||||
Date string `json:"date"`
|
||||
}
|
||||
|
||||
type queritGeoFilter struct {
|
||||
Countries *queritIncludeFilter `json:"countries,omitempty"`
|
||||
}
|
||||
|
||||
// QueritTool searches the web through Querit. Node-level parameters are
|
||||
// retained as defaults and model-emitted parameters override them per call.
|
||||
type QueritTool struct {
|
||||
helper *HTTPHelper
|
||||
envKey func() string
|
||||
defaults queritParams
|
||||
retryWait func(context.Context, int) bool
|
||||
}
|
||||
|
||||
var _ ToolComponent = (*QueritTool)(nil)
|
||||
var _ ReferenceBuilder = (*QueritTool)(nil)
|
||||
|
||||
// NewQueritTool returns a QueritTool using the shared HTTP helper and the
|
||||
// QUERIT_API_KEY environment variable.
|
||||
func NewQueritTool() *QueritTool {
|
||||
return newQueritTool(nil, nil, queritParams{}, nil)
|
||||
}
|
||||
|
||||
// NewQueritToolWith returns a QueritTool using the supplied HTTP helper.
|
||||
func NewQueritToolWith(helper *HTTPHelper) *QueritTool {
|
||||
return newQueritTool(helper, nil, queritParams{}, nil)
|
||||
}
|
||||
|
||||
// NewQueritToolWithEnvKey returns a QueritTool with an injectable API-key
|
||||
// resolver. It is intended for tests that must not depend on process state.
|
||||
func NewQueritToolWithEnvKey(helper *HTTPHelper, envKey func() string) *QueritTool {
|
||||
return newQueritTool(helper, envKey, queritParams{}, nil)
|
||||
}
|
||||
|
||||
func newQueritTool(
|
||||
helper *HTTPHelper,
|
||||
envKey func() string,
|
||||
defaults queritParams,
|
||||
retryWait func(context.Context, int) bool,
|
||||
) *QueritTool {
|
||||
if helper == nil {
|
||||
helper = NewHTTPHelper()
|
||||
}
|
||||
if envKey == nil {
|
||||
envKey = defaultQueritEnvKey
|
||||
}
|
||||
if defaults.Count == 0 {
|
||||
defaults.Count = 10
|
||||
}
|
||||
if defaults.ChunksPerDoc == nil {
|
||||
defaults.ChunksPerDoc = queritInt(3)
|
||||
}
|
||||
if retryWait == nil {
|
||||
retryWait = waitForQueritRetry
|
||||
}
|
||||
return &QueritTool{helper: helper, envKey: envKey, defaults: defaults, retryWait: retryWait}
|
||||
}
|
||||
|
||||
func defaultQueritEnvKey() string { return common.GetEnv(common.EnvQueritAPIKey) }
|
||||
|
||||
// Info returns the arguments that a chat model may provide. Credentials stay
|
||||
// in node configuration or the environment and are never exposed to the model.
|
||||
func (q *QueritTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
||||
return &schema.ToolInfo{
|
||||
Name: queritToolName,
|
||||
Desc: queritToolDescription,
|
||||
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
||||
"query": {
|
||||
Type: schema.String,
|
||||
Desc: "Search query.",
|
||||
Required: true,
|
||||
},
|
||||
"count": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Maximum number of documents to return. Defaults to 10 and must be at least 1.",
|
||||
Required: false,
|
||||
},
|
||||
"chunks_per_doc": {
|
||||
Type: schema.Integer,
|
||||
Desc: "Number of snippets per document. Defaults to 3 and must be between 1 and 3.",
|
||||
Required: false,
|
||||
},
|
||||
"site_include": {
|
||||
Type: schema.Array,
|
||||
ElemInfo: &schema.ParameterInfo{Type: schema.String},
|
||||
Desc: "Sites that search results must include.",
|
||||
Required: false,
|
||||
},
|
||||
"site_exclude": {
|
||||
Type: schema.Array,
|
||||
ElemInfo: &schema.ParameterInfo{Type: schema.String},
|
||||
Desc: "Sites that search results must exclude.",
|
||||
Required: false,
|
||||
},
|
||||
"time_range": {
|
||||
Type: schema.String,
|
||||
Desc: "Relative time range such as d7, w1, m3, or y1, or YYYY-MM-DDtoYYYY-MM-DD.",
|
||||
Required: false,
|
||||
},
|
||||
"country_include": {
|
||||
Type: schema.Array,
|
||||
ElemInfo: &schema.ParameterInfo{Type: schema.String},
|
||||
Desc: "Countries that search results must include.",
|
||||
Required: false,
|
||||
},
|
||||
"language_include": {
|
||||
Type: schema.Array,
|
||||
ElemInfo: &schema.ParameterInfo{Type: schema.String},
|
||||
Desc: "Languages that search results must include.",
|
||||
Required: false,
|
||||
},
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// InvokableRun performs a Querit search. All expected operational failures are
|
||||
// returned as soft-error JSON with a nil Go error so an Agent run can continue.
|
||||
func (q *QueritTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
||||
provided := make(map[string]json.RawMessage)
|
||||
if err := json.Unmarshal([]byte(argsJSON), &provided); err != nil {
|
||||
return queritErrorJSON(fmt.Errorf("parse arguments: %w", err)), nil
|
||||
}
|
||||
queryJSON, hasQuery := provided["query"]
|
||||
if hasQuery {
|
||||
if strings.TrimSpace(string(queryJSON)) == "null" {
|
||||
return queritErrorJSON(fmt.Errorf("query must be provided as a string")), nil
|
||||
}
|
||||
var query string
|
||||
if err := json.Unmarshal(queryJSON, &query); err != nil {
|
||||
return queritErrorJSON(fmt.Errorf("query must be provided as a string")), nil
|
||||
}
|
||||
}
|
||||
var runtimeParams queritParams
|
||||
if err := json.Unmarshal([]byte(argsJSON), &runtimeParams); err != nil {
|
||||
return queritErrorJSON(fmt.Errorf("parse arguments: %w", err)), nil
|
||||
}
|
||||
params := mergeQueritParams(q.defaults, runtimeParams, provided)
|
||||
if !hasQuery && strings.TrimSpace(params.Query) == "" {
|
||||
return queritErrorJSON(fmt.Errorf("query must be provided as a string")), nil
|
||||
}
|
||||
if params.Query == "" {
|
||||
return `{}`, nil
|
||||
}
|
||||
if err := validateQueritParams(params); err != nil {
|
||||
return queritErrorJSON(err, params.APIKey), nil
|
||||
}
|
||||
|
||||
apiKey := strings.TrimSpace(params.APIKey)
|
||||
if apiKey == "" {
|
||||
apiKey = strings.TrimSpace(q.envKey())
|
||||
}
|
||||
if apiKey == "" {
|
||||
return queritErrorJSON(fmt.Errorf("api_key is required (or set QUERIT_API_KEY)")), nil
|
||||
}
|
||||
|
||||
body, err := json.Marshal(buildQueritRequest(params))
|
||||
if err != nil {
|
||||
return queritErrorJSON(fmt.Errorf("encode request: %w", err), apiKey), nil
|
||||
}
|
||||
for attempt := 1; attempt <= queritMaxAttempts; attempt++ {
|
||||
resp, requestErr := q.helper.Do(
|
||||
ctx,
|
||||
http.MethodPost,
|
||||
queritEndpoint,
|
||||
string(body),
|
||||
"application/json",
|
||||
map[string]string{"Authorization": "Bearer " + apiKey},
|
||||
)
|
||||
if requestErr != nil {
|
||||
return queritErrorJSON(fmt.Errorf("request failed: %w", requestErr), apiKey), nil
|
||||
}
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if attempt == queritMaxAttempts || !q.retryWait(ctx, attempt) {
|
||||
return queritErrorJSON(fmt.Errorf("upstream returned %d after %d attempts", resp.StatusCode, attempt), apiKey), nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
return queritErrorJSON(fmt.Errorf("upstream returned %d", resp.StatusCode), apiKey), nil
|
||||
}
|
||||
|
||||
raw, readErr := io.ReadAll(resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
if readErr != nil {
|
||||
return queritErrorJSON(fmt.Errorf("read response: %w", readErr), apiKey), nil
|
||||
}
|
||||
if _, err := decodeQueritResponse(raw); err != nil {
|
||||
return queritErrorJSON(err, apiKey), nil
|
||||
}
|
||||
return string(raw), nil
|
||||
}
|
||||
|
||||
return queritErrorJSON(fmt.Errorf("request exhausted retries"), apiKey), nil
|
||||
}
|
||||
|
||||
func mergeQueritParams(defaults, runtimeParams queritParams, provided map[string]json.RawMessage) queritParams {
|
||||
merged := defaults
|
||||
if _, ok := provided["query"]; ok {
|
||||
merged.Query = runtimeParams.Query
|
||||
}
|
||||
if _, ok := provided["count"]; ok {
|
||||
merged.Count = runtimeParams.Count
|
||||
}
|
||||
if _, ok := provided["chunks_per_doc"]; ok {
|
||||
merged.ChunksPerDoc = runtimeParams.ChunksPerDoc
|
||||
}
|
||||
if _, ok := provided["site_include"]; ok {
|
||||
merged.SiteInclude = runtimeParams.SiteInclude
|
||||
}
|
||||
if _, ok := provided["site_exclude"]; ok {
|
||||
merged.SiteExclude = runtimeParams.SiteExclude
|
||||
}
|
||||
if _, ok := provided["time_range"]; ok {
|
||||
merged.TimeRange = runtimeParams.TimeRange
|
||||
}
|
||||
if _, ok := provided["country_include"]; ok {
|
||||
merged.CountryInclude = runtimeParams.CountryInclude
|
||||
}
|
||||
if _, ok := provided["language_include"]; ok {
|
||||
merged.LanguageInclude = runtimeParams.LanguageInclude
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
func validateQueritParams(params queritParams) error {
|
||||
if params.Count < 1 {
|
||||
return fmt.Errorf("count must be at least 1")
|
||||
}
|
||||
if params.ChunksPerDoc != nil && (*params.ChunksPerDoc < 1 || *params.ChunksPerDoc > 3) {
|
||||
return fmt.Errorf("chunks_per_doc must be between 1 and 3")
|
||||
}
|
||||
if !isValidQueritTimeRange(strings.TrimSpace(params.TimeRange)) {
|
||||
return fmt.Errorf("time_range must use dN, wN, mN, yN, or YYYY-MM-DDtoYYYY-MM-DD")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func isValidQueritTimeRange(value string) bool {
|
||||
return value == "" || queritRelativeTimeRangePattern.MatchString(value) || queritDateRangePattern.MatchString(value)
|
||||
}
|
||||
|
||||
func buildQueritRequest(params queritParams) queritRequest {
|
||||
request := queritRequest{
|
||||
Query: params.Query,
|
||||
Count: params.Count,
|
||||
ChunksPerDoc: params.ChunksPerDoc,
|
||||
}
|
||||
filters := &queritFilters{}
|
||||
if len(params.SiteInclude) > 0 || len(params.SiteExclude) > 0 {
|
||||
filters.Sites = &queritIncludeExcludeFilter{Include: params.SiteInclude, Exclude: params.SiteExclude}
|
||||
}
|
||||
if strings.TrimSpace(params.TimeRange) != "" {
|
||||
filters.TimeRange = &queritDateFilter{Date: params.TimeRange}
|
||||
}
|
||||
if len(params.CountryInclude) > 0 {
|
||||
filters.Geo = &queritGeoFilter{Countries: &queritIncludeFilter{Include: params.CountryInclude}}
|
||||
}
|
||||
if len(params.LanguageInclude) > 0 {
|
||||
filters.Languages = &queritIncludeFilter{Include: params.LanguageInclude}
|
||||
}
|
||||
if filters.Sites != nil || filters.TimeRange != nil || filters.Geo != nil || filters.Languages != nil {
|
||||
request.Filters = filters
|
||||
}
|
||||
return request
|
||||
}
|
||||
|
||||
func decodeQueritResponse(raw []byte) (map[string]any, error) {
|
||||
var decoded any
|
||||
decoder := json.NewDecoder(strings.NewReader(string(raw)))
|
||||
decoder.UseNumber()
|
||||
if err := decoder.Decode(&decoded); err != nil {
|
||||
return nil, fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if err := decoder.Decode(&struct{}{}); err != io.EOF {
|
||||
if err == nil {
|
||||
return nil, fmt.Errorf("decode response: expected exactly one JSON value")
|
||||
}
|
||||
return nil, fmt.Errorf("decode response: trailing content: %w", err)
|
||||
}
|
||||
response, ok := decoded.(map[string]any)
|
||||
if !ok || response == nil {
|
||||
return nil, fmt.Errorf("decode response: expected a JSON object")
|
||||
}
|
||||
resultsValue, exists := response["results"]
|
||||
if !exists {
|
||||
return response, nil
|
||||
}
|
||||
results, ok := resultsValue.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("decode response: results must be a JSON object")
|
||||
}
|
||||
resultValue, exists := results["result"]
|
||||
if !exists {
|
||||
return response, nil
|
||||
}
|
||||
if _, ok := resultValue.([]any); !ok {
|
||||
return nil, fmt.Errorf("decode response: results.result must be a JSON array")
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func waitForQueritRetry(ctx context.Context, attempt int) bool {
|
||||
delay := 200 * time.Millisecond
|
||||
for current := 1; current < attempt; current++ {
|
||||
delay *= 2
|
||||
}
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case <-timer.C:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// ComponentSpec returns the QueritSearch Canvas-facing metadata.
|
||||
func (q *QueritTool) ComponentSpec() ComponentSpec {
|
||||
return ComponentSpec{
|
||||
PreserveJSONNumbers: true,
|
||||
Inputs: map[string]string{
|
||||
"api_key": "Querit API key. Uses QUERIT_API_KEY when empty.",
|
||||
"query": "Search query.",
|
||||
"count": "Maximum number of documents.",
|
||||
"chunks_per_doc": "Number of snippets per document.",
|
||||
"site_include": "Sites that search results must include.",
|
||||
"site_exclude": "Sites that search results must exclude.",
|
||||
"time_range": "dN, wN, mN, yN, or YYYY-MM-DDtoYYYY-MM-DD.",
|
||||
"country_include": "Countries that search results must include.",
|
||||
"language_include": "Languages that search results must include.",
|
||||
},
|
||||
Outputs: map[string]string{
|
||||
"formalized_content": "Rendered Querit references for downstream prompts.",
|
||||
"json": "Complete raw Querit JSON response.",
|
||||
},
|
||||
InputForm: map[string]any{
|
||||
"api_key": map[string]any{"name": "API key", "type": "password"},
|
||||
"query": map[string]any{"name": "Query", "type": "line"},
|
||||
"count": map[string]any{"name": "Count", "type": "line"},
|
||||
"chunks_per_doc": map[string]any{"name": "Chunks per document", "type": "line"},
|
||||
"site_include": map[string]any{"name": "Site include", "type": "line"},
|
||||
"site_exclude": map[string]any{"name": "Site exclude", "type": "line"},
|
||||
"time_range": map[string]any{"name": "Time range", "type": "line"},
|
||||
"country_include": map[string]any{"name": "Country include", "type": "line"},
|
||||
"language_include": map[string]any{"name": "Language include", "type": "line"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// BuildReferences converts results.result into RAGFlow retrieval references.
|
||||
func (q *QueritTool) BuildReferences(_ context.Context, response map[string]any) ([]map[string]any, []map[string]any) {
|
||||
results := queritResults(response)
|
||||
chunks := make([]map[string]any, 0, len(results))
|
||||
docAggs := make([]map[string]any, 0, len(results))
|
||||
for _, result := range results {
|
||||
title := queritText(result["title"])
|
||||
resultURL := queritText(result["url"])
|
||||
content := queritText(result["snippet"])
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
content = queritDataImagePattern.ReplaceAllString(content, "")
|
||||
content = truncateQueritRunes(content, 10000)
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
documentID := strconv.FormatInt(queritHashInt(content, 100000000), 10)
|
||||
displayID := strconv.FormatInt(queritHashInt(documentID, 500), 10)
|
||||
chunks = append(chunks, map[string]any{
|
||||
"id": displayID,
|
||||
"chunk_id": documentID,
|
||||
"content": content,
|
||||
"doc_id": documentID,
|
||||
"document_id": documentID,
|
||||
"docnm_kwd": title,
|
||||
"document_name": title,
|
||||
"similarity": 1,
|
||||
"score": 1,
|
||||
"url": resultURL,
|
||||
})
|
||||
docAggs = append(docAggs, map[string]any{
|
||||
"doc_name": title,
|
||||
"doc_id": documentID,
|
||||
"count": 1,
|
||||
"url": resultURL,
|
||||
})
|
||||
}
|
||||
return chunks, docAggs
|
||||
}
|
||||
|
||||
// BuildComponentOutputs keeps the complete upstream object in json and adds a
|
||||
// reference-formatted string for downstream Canvas components.
|
||||
func (q *QueritTool) BuildComponentOutputs(response map[string]any) map[string]any {
|
||||
chunks, _ := q.BuildReferences(context.Background(), response)
|
||||
return map[string]any{
|
||||
"formalized_content": renderQueritReferences(chunks, queritPromptMaxTokens),
|
||||
"json": response,
|
||||
}
|
||||
}
|
||||
|
||||
func queritResults(response map[string]any) []map[string]any {
|
||||
results, ok := response["results"].(map[string]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
raw, ok := results["result"].([]any)
|
||||
if !ok {
|
||||
if typed, typedOK := results["result"].([]map[string]any); typedOK {
|
||||
return typed
|
||||
}
|
||||
return nil
|
||||
}
|
||||
items := make([]map[string]any, 0, len(raw))
|
||||
for _, value := range raw {
|
||||
if item, itemOK := value.(map[string]any); itemOK {
|
||||
items = append(items, item)
|
||||
}
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
func renderQueritReferences(chunks []map[string]any, maxTokens int) string {
|
||||
usedTokens := 0
|
||||
blocks := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
content := queritText(chunk["content"])
|
||||
usedTokens += tokenizer.NumTokensFromString(content)
|
||||
blocks = append(blocks, strings.Join([]string{
|
||||
"\nID: " + queritText(chunk["id"]),
|
||||
"├── Title: " + queritPromptField(chunk["document_name"]),
|
||||
"├── URL: " + queritPromptField(chunk["url"]),
|
||||
"└── Content:\n" + content,
|
||||
}, "\n"))
|
||||
if maxTokens > 0 && float64(maxTokens)*0.97 < float64(usedTokens) {
|
||||
break
|
||||
}
|
||||
}
|
||||
return strings.Join(blocks, "\n")
|
||||
}
|
||||
|
||||
func queritText(value any) string {
|
||||
if value == nil {
|
||||
return ""
|
||||
}
|
||||
if text, ok := value.(string); ok {
|
||||
return text
|
||||
}
|
||||
return fmt.Sprint(value)
|
||||
}
|
||||
|
||||
func queritPromptField(value any) string {
|
||||
return queritNewlinePattern.ReplaceAllString(queritText(value), " ")
|
||||
}
|
||||
|
||||
func queritHashInt(value string, modulus int64) int64 {
|
||||
digest := sha1.Sum([]byte(value))
|
||||
number := new(big.Int).SetBytes(digest[:])
|
||||
return new(big.Int).Mod(number, big.NewInt(modulus)).Int64()
|
||||
}
|
||||
|
||||
func truncateQueritRunes(value string, limit int) string {
|
||||
if limit <= 0 {
|
||||
return ""
|
||||
}
|
||||
runes := []rune(value)
|
||||
if len(runes) <= limit {
|
||||
return value
|
||||
}
|
||||
return string(runes[:limit])
|
||||
}
|
||||
|
||||
func queritInt(value int) *int { return &value }
|
||||
|
||||
func queritStringSlice(value any) ([]string, bool) {
|
||||
switch items := value.(type) {
|
||||
case []string:
|
||||
return append([]string(nil), items...), true
|
||||
case []any:
|
||||
result := make([]string, 0, len(items))
|
||||
for _, item := range items {
|
||||
text, ok := item.(string)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
result = append(result, text)
|
||||
}
|
||||
return result, true
|
||||
case nil:
|
||||
return nil, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func queritErrorJSON(err error, apiKeys ...string) string {
|
||||
message := "querit_search: unknown error"
|
||||
if err != nil {
|
||||
message = "querit_search: " + err.Error()
|
||||
}
|
||||
for _, apiKey := range apiKeys {
|
||||
if apiKey = strings.TrimSpace(apiKey); apiKey != "" {
|
||||
message = strings.ReplaceAll(message, apiKey, "[REDACTED]")
|
||||
}
|
||||
}
|
||||
raw, marshalErr := json.Marshal(map[string]any{"_ERROR": message})
|
||||
if marshalErr != nil {
|
||||
return `{"_ERROR":"querit_search: marshal error"}`
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
594
internal/agent/tool/querit_test.go
Normal file
594
internal/agent/tool/querit_test.go
Normal file
@@ -0,0 +1,594 @@
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
|
||||
package tool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestQueritBuildsMinimalRequest(t *testing.T) {
|
||||
var gotMethod, gotPath, gotAuthorization, gotContentType string
|
||||
var gotBody map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
gotMethod = request.Method
|
||||
gotPath = request.URL.Path
|
||||
gotAuthorization = request.Header.Get("Authorization")
|
||||
gotContentType = request.Header.Get("Content-Type")
|
||||
_ = json.NewDecoder(request.Body).Decode(&gotBody)
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[]},"search_id":"search-1"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := newQueritTool(helper, func() string { return "" }, queritParams{APIKey: "key-test"}, nil)
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"ragflow"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if gotMethod != http.MethodPost || gotPath != "/v1/search" {
|
||||
t.Fatalf("request = %s %s, want POST /v1/search", gotMethod, gotPath)
|
||||
}
|
||||
if gotAuthorization != "Bearer key-test" {
|
||||
t.Fatalf("Authorization = %q", gotAuthorization)
|
||||
}
|
||||
if !strings.HasPrefix(gotContentType, "application/json") {
|
||||
t.Fatalf("Content-Type = %q", gotContentType)
|
||||
}
|
||||
if gotBody["query"] != "ragflow" || gotBody["count"] != float64(10) || gotBody["chunksPerDoc"] != float64(3) {
|
||||
t.Fatalf("request body = %#v", gotBody)
|
||||
}
|
||||
if _, exists := gotBody["filters"]; exists {
|
||||
t.Fatalf("empty filters must be omitted: %#v", gotBody)
|
||||
}
|
||||
if !strings.Contains(out, `"search_id":"search-1"`) {
|
||||
t.Fatalf("complete response was not retained: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritBuildsFiltersAndMergesRuntimeOverrides(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
_ = json.NewDecoder(request.Body).Decode(&gotBody)
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
defaults := queritParams{
|
||||
APIKey: "stored-key",
|
||||
Count: 20,
|
||||
ChunksPerDoc: queritInt(2),
|
||||
SiteInclude: []string{"stored.example"},
|
||||
SiteExclude: []string{"blocked.example"},
|
||||
TimeRange: "w1",
|
||||
CountryInclude: []string{"CN"},
|
||||
LanguageInclude: []string{"zh"},
|
||||
}
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := newQueritTool(helper, func() string { return "" }, defaults, func(context.Context, int) bool { return true })
|
||||
_, err := querit.InvokableRun(context.Background(), `{"query":"ragflow","count":5,"site_include":[],"language_include":["en"]}`)
|
||||
if err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if gotBody["count"] != float64(5) || gotBody["chunksPerDoc"] != float64(2) {
|
||||
t.Fatalf("merged scalar defaults = %#v", gotBody)
|
||||
}
|
||||
filters := gotBody["filters"].(map[string]any)
|
||||
sites := filters["sites"].(map[string]any)
|
||||
if _, exists := sites["include"]; exists || len(sites["exclude"].([]any)) != 1 {
|
||||
t.Fatalf("explicit empty site_include did not clear node default: %#v", sites)
|
||||
}
|
||||
if filters["timeRange"].(map[string]any)["date"] != "w1" {
|
||||
t.Fatalf("timeRange = %#v", filters["timeRange"])
|
||||
}
|
||||
if filters["geo"].(map[string]any)["countries"].(map[string]any)["include"].([]any)[0] != "CN" {
|
||||
t.Fatalf("geo filter = %#v", filters["geo"])
|
||||
}
|
||||
if filters["languages"].(map[string]any)["include"].([]any)[0] != "en" {
|
||||
t.Fatalf("language filter = %#v", filters["languages"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritUsesNodeQueryWhenRuntimeQueryIsOmitted(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
_ = json.NewDecoder(request.Body).Decode(&gotBody)
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := newQueritTool(
|
||||
helper,
|
||||
func() string { return "" },
|
||||
queritParams{APIKey: "stored-key", Query: "node query"},
|
||||
nil,
|
||||
)
|
||||
out, err := querit.InvokableRun(context.Background(), `{}`)
|
||||
if err != nil || strings.Contains(out, "_ERROR") {
|
||||
t.Fatalf("InvokableRun = %s, %v", out, err)
|
||||
}
|
||||
if gotBody["query"] != "node query" {
|
||||
t.Fatalf("request query = %#v, want node query", gotBody["query"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritAPIKeyResolutionAndEmptyQuery(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
var authorization string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
calls.Add(1)
|
||||
authorization = request.Header.Get("Authorization")
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return "environment-secret" })
|
||||
if _, err := querit.InvokableRun(context.Background(), `{"query":"ragflow"}`); err != nil {
|
||||
t.Fatalf("InvokableRun: %v", err)
|
||||
}
|
||||
if authorization != "Bearer environment-secret" {
|
||||
t.Fatalf("Authorization = %q", authorization)
|
||||
}
|
||||
emptyOut, err := querit.InvokableRun(context.Background(), `{"query":""}`)
|
||||
if err != nil {
|
||||
t.Fatalf("empty query: %v", err)
|
||||
}
|
||||
if calls.Load() != 1 || emptyOut != `{}` {
|
||||
t.Fatalf("empty query result = %s; calls = %d", emptyOut, calls.Load())
|
||||
}
|
||||
|
||||
missing := NewQueritToolWithEnvKey(helper, func() string { return "" })
|
||||
out, err := missing.InvokableRun(context.Background(), `{"query":"ragflow"}`)
|
||||
if err != nil {
|
||||
t.Fatalf("missing key returned Go error: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, "api_key") || strings.Contains(out, "environment-secret") || calls.Load() != 1 {
|
||||
t.Fatalf("missing-key result = %s, calls = %d", out, calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritRuntimeAPIKeyCannotOverrideNodeConfiguration(t *testing.T) {
|
||||
var authorization string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
authorization = request.Header.Get("Authorization")
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := newQueritTool(helper, func() string { return "" }, queritParams{APIKey: "stored-key"}, nil)
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"ragflow","api_key":"runtime-key"}`)
|
||||
if err != nil || strings.Contains(out, "_ERROR") {
|
||||
t.Fatalf("InvokableRun = %s, %v", out, err)
|
||||
}
|
||||
if authorization != "Bearer stored-key" {
|
||||
t.Fatalf("Authorization = %q, want stored node key", authorization)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritRejectsMissingNullAndNonStringQueries(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: roundTripperErrorFunc(func(*http.Request) error {
|
||||
calls.Add(1)
|
||||
return errors.New("network must not be called")
|
||||
})})
|
||||
querit := NewQueritToolWith(helper)
|
||||
for _, args := range []string{`{}`, `{"query":null}`, `{"query":123}`} {
|
||||
t.Run(args, func(t *testing.T) {
|
||||
out, err := querit.InvokableRun(context.Background(), args)
|
||||
if err != nil || !strings.Contains(out, "_ERROR") || !strings.Contains(out, "query") {
|
||||
t.Fatalf("InvokableRun(%s) = %s, %v", args, out, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
if calls.Load() != 0 {
|
||||
t.Fatalf("invalid queries made %d network calls", calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritExplicitNullChunksPerDocIsOmitted(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
_ = json.NewDecoder(request.Body).Decode(&gotBody)
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return "k" })
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x","chunks_per_doc":null}`)
|
||||
if err != nil || strings.Contains(out, "_ERROR") {
|
||||
t.Fatalf("InvokableRun = %s, %v", out, err)
|
||||
}
|
||||
if _, exists := gotBody["chunksPerDoc"]; exists {
|
||||
t.Fatalf("explicit null chunks_per_doc was not omitted: %#v", gotBody)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritValidatesParametersBeforeRequest(t *testing.T) {
|
||||
tests := []string{
|
||||
`{"query":"x","count":0}`,
|
||||
`{"query":"x","chunks_per_doc":4}`,
|
||||
`{"query":"x","time_range":"last week"}`,
|
||||
`{"query":"x","time_range":"2026-01-01,2026-01-31"}`,
|
||||
`{"query":"x","site_include":[1]}`,
|
||||
}
|
||||
for _, args := range tests {
|
||||
t.Run(args, func(t *testing.T) {
|
||||
out, err := NewQueritTool().InvokableRun(context.Background(), args)
|
||||
if err != nil || !strings.Contains(out, "_ERROR") {
|
||||
t.Fatalf("InvokableRun(%s) = %s, %v", args, out, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritTimeRangeContract(t *testing.T) {
|
||||
tests := []struct {
|
||||
value string
|
||||
valid bool
|
||||
}{
|
||||
{value: "", valid: true},
|
||||
{value: "d7", valid: true},
|
||||
{value: "w2", valid: true},
|
||||
{value: "m3", valid: true},
|
||||
{value: "y1", valid: true},
|
||||
{value: "2026-01-01to2026-01-31", valid: true},
|
||||
{value: "d0", valid: false},
|
||||
{value: "7d", valid: false},
|
||||
{value: "2026-01-01,2026-01-31", valid: false},
|
||||
{value: "2026-01-01..2026-01-31", valid: false},
|
||||
{value: "2026-01-01-2026-01-31", valid: false},
|
||||
{value: "last week", valid: false},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.value, func(t *testing.T) {
|
||||
if got := isValidQueritTimeRange(test.value); got != test.valid {
|
||||
t.Fatalf("isValidQueritTimeRange(%q) = %v, want %v", test.value, got, test.valid)
|
||||
}
|
||||
if !test.valid || test.value == "" {
|
||||
return
|
||||
}
|
||||
request := buildQueritRequest(queritParams{TimeRange: test.value})
|
||||
if request.Filters == nil || request.Filters.TimeRange == nil || request.Filters.TimeRange.Date != test.value {
|
||||
t.Fatalf("time range request mapping = %#v", request.Filters)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritHTTPFailuresAreSoftErrors(t *testing.T) {
|
||||
t.Run("unauthorized is not retried", func(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
calls.Add(1)
|
||||
writer.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = writer.Write([]byte(`{"message":"do not expose upstream bodies"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return "secret-key" })
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || calls.Load() != 1 || !strings.Contains(out, "401") || strings.Contains(out, "secret-key") {
|
||||
t.Fatalf("result = %s, err = %v, calls = %d", out, err, calls.Load())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("rate limit is retried at most three times", func(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
calls.Add(1)
|
||||
writer.WriteHeader(http.StatusTooManyRequests)
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := newQueritTool(helper, nil, queritParams{APIKey: "secret-key"}, func(context.Context, int) bool { return true })
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || calls.Load() != queritMaxAttempts || !strings.Contains(out, "429") || strings.Contains(out, "secret-key") {
|
||||
t.Fatalf("result = %s, err = %v, calls = %d", out, err, calls.Load())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("HTTP helper retries temporary server errors", func(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
attempt := calls.Add(1)
|
||||
if attempt < 3 {
|
||||
writer.WriteHeader(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
_, _ = writer.Write([]byte(`{"results":{"result":[]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelperWithRetry(RetryConfig{
|
||||
MaxAttempts: 3,
|
||||
BaseBackoff: time.Nanosecond,
|
||||
MaxBackoff: time.Nanosecond,
|
||||
}).WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
out, err := NewQueritToolWithEnvKey(helper, func() string { return "k" }).InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || calls.Load() != 3 || strings.Contains(out, "_ERROR") {
|
||||
t.Fatalf("result = %s, err = %v, calls = %d", out, err, calls.Load())
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("persistent server errors are soft", func(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
calls.Add(1)
|
||||
writer.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelperWithRetry(RetryConfig{
|
||||
MaxAttempts: 3,
|
||||
BaseBackoff: time.Nanosecond,
|
||||
MaxBackoff: time.Nanosecond,
|
||||
}).WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return "environment-secret" })
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || calls.Load() != 3 || !strings.Contains(out, "_ERROR") || !strings.Contains(out, "500") {
|
||||
t.Fatalf("result = %s, err = %v, calls = %d", out, err, calls.Load())
|
||||
}
|
||||
})
|
||||
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
secret string
|
||||
node bool
|
||||
}{
|
||||
{name: "node API key is redacted from transport errors", secret: "node-secret", node: true},
|
||||
{name: "environment API key is redacted from transport errors", secret: "environment-secret"},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
helper := NewHTTPHelperWithRetry(RetryConfig{
|
||||
MaxAttempts: 1,
|
||||
BaseBackoff: time.Nanosecond,
|
||||
MaxBackoff: time.Nanosecond,
|
||||
}).WithClient(&http.Client{Transport: roundTripperErrorFunc(func(*http.Request) error {
|
||||
return fmt.Errorf("transport rejected Bearer %s", test.secret)
|
||||
})})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return test.secret })
|
||||
if test.node {
|
||||
querit = newQueritTool(helper, func() string { return "" }, queritParams{APIKey: test.secret}, nil)
|
||||
}
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || !strings.Contains(out, "_ERROR") || !strings.Contains(out, "[REDACTED]") {
|
||||
t.Fatalf("result = %s, err = %v", out, err)
|
||||
}
|
||||
if strings.Contains(out, test.secret) {
|
||||
t.Fatalf("soft error exposed API key: %s", out)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("network errors are soft", func(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
helper := NewHTTPHelperWithRetry(RetryConfig{
|
||||
MaxAttempts: 3,
|
||||
BaseBackoff: time.Nanosecond,
|
||||
MaxBackoff: time.Nanosecond,
|
||||
}).WithClient(&http.Client{Transport: roundTripperErrorFunc(func(*http.Request) error {
|
||||
calls.Add(1)
|
||||
return errors.New("offline")
|
||||
})})
|
||||
out, err := NewQueritToolWithEnvKey(helper, func() string { return "k" }).InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || calls.Load() != 3 || !strings.Contains(out, "_ERROR") {
|
||||
t.Fatalf("result = %s, err = %v, calls = %d", out, err, calls.Load())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestQueritRejectsInvalidJSONResponse(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte(`not-json`))
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return "k" })
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || !strings.Contains(out, "decode response") {
|
||||
t.Fatalf("result = %s, err = %v", out, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritRejectsMalformedResponseShapes(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
want string
|
||||
}{
|
||||
{name: "top-level array", body: `[]`, want: "JSON object"},
|
||||
{name: "results is not an object", body: `{"results":[]}`, want: "results must be a JSON object"},
|
||||
{name: "results is null", body: `{"results":null}`, want: "results must be a JSON object"},
|
||||
{name: "result is not an array", body: `{"results":{"result":{}}}`, want: "results.result must be a JSON array"},
|
||||
{name: "result is null", body: `{"results":{"result":null}}`, want: "results.result must be a JSON array"},
|
||||
{name: "trailing content", body: `{"results":{"result":[]}} trailing`, want: "trailing content"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte(test.body))
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return "k" })
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || !strings.Contains(out, "_ERROR") || !strings.Contains(out, test.want) {
|
||||
t.Fatalf("result = %s, err = %v", out, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritAcceptsMissingResultContainers(t *testing.T) {
|
||||
for _, body := range []string{`{}`, `{"results":{}}`} {
|
||||
t.Run(body, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
_, _ = writer.Write([]byte(body))
|
||||
}))
|
||||
defer server.Close()
|
||||
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteQueritHostTransport(server.URL)})
|
||||
querit := NewQueritToolWithEnvKey(helper, func() string { return "k" })
|
||||
out, err := querit.InvokableRun(context.Background(), `{"query":"x"}`)
|
||||
if err != nil || strings.Contains(out, "_ERROR") {
|
||||
t.Fatalf("result = %s, err = %v", out, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritReferencesAndCompleteComponentOutput(t *testing.T) {
|
||||
response := map[string]any{
|
||||
"search_id": "search-1",
|
||||
"query_context": map[string]any{"rewritten": "rag flow"},
|
||||
"results": map[string]any{"result": []any{
|
||||
map[string]any{"title": "RAGFlow", "url": "https://ragflow.io", "snippet": "RAG engine", "extra": true},
|
||||
map[string]any{"title": "Missing optional values"},
|
||||
"invalid item",
|
||||
}},
|
||||
}
|
||||
querit := NewQueritTool()
|
||||
chunks, docAggs := querit.BuildReferences(context.Background(), response)
|
||||
if len(chunks) != 1 || len(docAggs) != 1 {
|
||||
t.Fatalf("references = %#v / %#v", chunks, docAggs)
|
||||
}
|
||||
if chunks[0]["content"] != "RAG engine" || chunks[0]["score"] != 1 || chunks[0]["similarity"] != 1 {
|
||||
t.Fatalf("reference = %#v", chunks[0])
|
||||
}
|
||||
outputs := querit.BuildComponentOutputs(response)
|
||||
if outputs["json"].(map[string]any)["search_id"] != "search-1" {
|
||||
t.Fatalf("complete json output = %#v", outputs["json"])
|
||||
}
|
||||
formatted := outputs["formalized_content"].(string)
|
||||
for _, expected := range []string{"Title: RAGFlow", "URL: https://ragflow.io", "RAG engine"} {
|
||||
if !strings.Contains(formatted, expected) {
|
||||
t.Fatalf("formalized_content missing %q: %s", expected, formatted)
|
||||
}
|
||||
}
|
||||
if chunks, docAggs := querit.BuildReferences(context.Background(), map[string]any{"results": nil}); len(chunks) != 0 || len(docAggs) != 0 {
|
||||
t.Fatalf("malformed response references = %#v / %#v", chunks, docAggs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritReferencesSanitizeAndLimitSnippets(t *testing.T) {
|
||||
longSnippet := strings.Repeat("界", 10001)
|
||||
response := map[string]any{"results": map[string]any{"result": []any{
|
||||
map[string]any{"title": "empty", "snippet": ""},
|
||||
map[string]any{"title": "image only", "snippet": ""},
|
||||
map[string]any{"title": "cleaned", "snippet": "before  after"},
|
||||
map[string]any{"title": "limited", "snippet": longSnippet},
|
||||
}}}
|
||||
chunks, docAggs := NewQueritTool().BuildReferences(context.Background(), response)
|
||||
if len(chunks) != 2 || len(docAggs) != 2 {
|
||||
t.Fatalf("references = %#v / %#v", chunks, docAggs)
|
||||
}
|
||||
if chunks[0]["content"] != "before after" {
|
||||
t.Fatalf("base64 image was not removed: %#v", chunks[0])
|
||||
}
|
||||
limited, _ := chunks[1]["content"].(string)
|
||||
if len([]rune(limited)) != 10000 {
|
||||
t.Fatalf("limited snippet length = %d", len([]rune(limited)))
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueritInfoDoesNotExposeAPIKey(t *testing.T) {
|
||||
info, err := NewQueritTool().Info(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("Info: %v", err)
|
||||
}
|
||||
if info.Name != queritToolName || info.Desc == "" || info.ParamsOneOf == nil {
|
||||
t.Fatalf("Info = %#v", info)
|
||||
}
|
||||
encoded, err := json.Marshal(info)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal Info: %v", err)
|
||||
}
|
||||
if strings.Contains(string(encoded), "api_key") {
|
||||
t.Fatalf("Info exposed API key: %s", encoded)
|
||||
}
|
||||
jsonSchema, err := info.ParamsOneOf.ToJSONSchema()
|
||||
if err != nil {
|
||||
t.Fatalf("ToJSONSchema: %v", err)
|
||||
}
|
||||
rawSchema, err := json.Marshal(jsonSchema)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal schema: %v", err)
|
||||
}
|
||||
var paramsSchema map[string]any
|
||||
if err := json.Unmarshal(rawSchema, ¶msSchema); err != nil {
|
||||
t.Fatalf("decode schema: %v", err)
|
||||
}
|
||||
properties, ok := paramsSchema["properties"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("schema properties = %#v", paramsSchema["properties"])
|
||||
}
|
||||
for _, name := range []string{"site_include", "site_exclude", "country_include", "language_include"} {
|
||||
property, ok := properties[name].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("%s schema = %#v", name, properties[name])
|
||||
}
|
||||
items, ok := property["items"].(map[string]any)
|
||||
if !ok || items["type"] != "string" {
|
||||
t.Fatalf("%s items schema = %#v", name, property["items"])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type roundTripperErrorFunc func(*http.Request) error
|
||||
|
||||
func (f roundTripperErrorFunc) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
return nil, f(request)
|
||||
}
|
||||
|
||||
func rewriteQueritHostTransport(serverURL string) http.RoundTripper {
|
||||
target, err := url.Parse(serverURL)
|
||||
if err != nil {
|
||||
panic("rewriteQueritHostTransport: bad server URL: " + err.Error())
|
||||
}
|
||||
return &queritHostSwapTransport{
|
||||
inner: http.DefaultTransport,
|
||||
host: target.Host,
|
||||
scheme: target.Scheme,
|
||||
}
|
||||
}
|
||||
|
||||
type queritHostSwapTransport struct {
|
||||
inner http.RoundTripper
|
||||
host string
|
||||
scheme string
|
||||
}
|
||||
|
||||
func (transport *queritHostSwapTransport) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
cloned := request.Clone(request.Context())
|
||||
cloned.URL.Scheme = transport.scheme
|
||||
cloned.URL.Host = transport.host
|
||||
cloned.Host = transport.host
|
||||
return transport.inner.RoundTrip(cloned)
|
||||
}
|
||||
@@ -49,6 +49,9 @@ var registry = map[string]Factory{
|
||||
"keenable": buildKeenableTool,
|
||||
"pubmed": buildPubMedTool,
|
||||
"qweather": noConfig("qweather", func() einotool.BaseTool { return NewQWeatherTool() }),
|
||||
"querit": buildQueritTool,
|
||||
"querit_search": buildQueritTool,
|
||||
"queritsearch": buildQueritTool,
|
||||
"retrieval": buildRetrievalTool,
|
||||
"search_my_dataset": buildRetrievalTool,
|
||||
"search_my_dateset": buildRetrievalTool,
|
||||
@@ -533,6 +536,61 @@ func buildTavilyTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
return newTavilyTool(nil, nil, defaults), nil
|
||||
}
|
||||
|
||||
func buildQueritTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
defaults := queritParams{}
|
||||
stringFields := map[string]*string{
|
||||
"api_key": &defaults.APIKey,
|
||||
"query": &defaults.Query,
|
||||
"time_range": &defaults.TimeRange,
|
||||
}
|
||||
for key, destination := range stringFields {
|
||||
value, exists := params[key]
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
text, valid := value.(string)
|
||||
if !valid {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires string node-level param %s", "querit_search", key)
|
||||
}
|
||||
*destination = text
|
||||
}
|
||||
if value, exists := params["count"]; exists {
|
||||
count, valid := strictInt(value)
|
||||
if !valid || count < 1 {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires integer node-level param count of at least 1", "querit_search")
|
||||
}
|
||||
defaults.Count = count
|
||||
}
|
||||
if value, exists := params["chunks_per_doc"]; exists {
|
||||
chunksPerDoc, valid := strictInt(value)
|
||||
if !valid || chunksPerDoc < 1 || chunksPerDoc > 3 {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires integer node-level param chunks_per_doc within [1, 3]", "querit_search")
|
||||
}
|
||||
defaults.ChunksPerDoc = queritInt(chunksPerDoc)
|
||||
}
|
||||
listFields := map[string]*[]string{
|
||||
"site_include": &defaults.SiteInclude,
|
||||
"site_exclude": &defaults.SiteExclude,
|
||||
"country_include": &defaults.CountryInclude,
|
||||
"language_include": &defaults.LanguageInclude,
|
||||
}
|
||||
for key, destination := range listFields {
|
||||
value, exists := params[key]
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
items, valid := queritStringSlice(value)
|
||||
if !valid {
|
||||
return nil, fmt.Errorf("agent tool: tool %q requires string array node-level param %s", "querit_search", key)
|
||||
}
|
||||
*destination = items
|
||||
}
|
||||
if !isValidQueritTimeRange(strings.TrimSpace(defaults.TimeRange)) {
|
||||
return nil, fmt.Errorf("agent tool: tool %q has invalid node-level param time_range", "querit_search")
|
||||
}
|
||||
return newQueritTool(nil, nil, defaults, nil), nil
|
||||
}
|
||||
|
||||
func buildKeenableTool(params map[string]any) (einotool.BaseTool, error) {
|
||||
defaults := keenableParams{}
|
||||
apiKey := ""
|
||||
|
||||
@@ -64,12 +64,43 @@ func TestBuildByName_TavilyCanvasComponentNames(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildByName_QueritAliases(t *testing.T) {
|
||||
for _, name := range []string{"querit", "querit_search", "queritsearch", "QueritSearch"} {
|
||||
built, err := BuildByName(name, map[string]any{
|
||||
"api_key": "stored-key",
|
||||
"count": float64(8),
|
||||
"chunks_per_doc": float64(2),
|
||||
"site_include": []any{"example.com"},
|
||||
"site_exclude": []string{"blocked.example"},
|
||||
"time_range": "d7",
|
||||
"country_include": []any{"CN"},
|
||||
"language_include": []any{"zh"},
|
||||
"outputs": map[string]any{"json": map[string]any{}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BuildByName(%q): %v", name, err)
|
||||
}
|
||||
querit, ok := built.(*QueritTool)
|
||||
if !ok {
|
||||
t.Fatalf("BuildByName(%q) returned %T, want *QueritTool", name, built)
|
||||
}
|
||||
if querit.defaults.Count != 8 || querit.defaults.ChunksPerDoc == nil || *querit.defaults.ChunksPerDoc != 2 || len(querit.defaults.SiteInclude) != 1 {
|
||||
t.Fatalf("BuildByName(%q) defaults = %#v", name, querit.defaults)
|
||||
}
|
||||
info, infoErr := built.Info(context.Background())
|
||||
if infoErr != nil || info.Name != "querit_search" {
|
||||
t.Fatalf("BuildByName(%q).Info() = %#v, %v", name, info, infoErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAll_AllRegisteredTools(t *testing.T) {
|
||||
// Every key in registry.
|
||||
names := []string{
|
||||
"akshare", "arxiv", "bgpt", "code_exec", "crawler", "deepl",
|
||||
"duckduckgo", "email", "exesql", "execute_sql", "github", "google",
|
||||
"google_scholar", "google_scholar_search", "jin10", "keenable", "pubmed", "qweather",
|
||||
"querit", "querit_search", "queritsearch",
|
||||
"retrieval", "search_my_dataset", "search_my_dateset", "searxng",
|
||||
"tavily", "tavily_extract", "tushare", "web_crawler", "wencai", "wikipedia", "wikipedia_search",
|
||||
"yahoo_finance",
|
||||
@@ -133,6 +164,7 @@ func TestToolRegistry_SchemasAreComplete(t *testing.T) {
|
||||
"akshare", "arxiv", "bgpt", "code_exec", "crawler", "deepl",
|
||||
"duckduckgo", "email", "execute_sql", "exesql", "github", "google",
|
||||
"google_scholar", "google_scholar_search", "jin10", "keenable", "pubmed", "qweather",
|
||||
"querit", "querit_search", "queritsearch",
|
||||
"retrieval", "search_my_dataset", "search_my_dateset", "searxng",
|
||||
"tavily", "tavily_extract", "tushare", "web_crawler", "wencai", "wikipedia", "wikipedia_search",
|
||||
"yahoo_finance",
|
||||
@@ -202,6 +234,9 @@ func TestToolRegistry_SchemasAreComplete(t *testing.T) {
|
||||
"web_crawler": "web_crawler",
|
||||
"wikipedia": "wikipedia_search",
|
||||
"wikipedia_search": "wikipedia_search",
|
||||
"querit": "querit_search",
|
||||
"querit_search": "querit_search",
|
||||
"queritsearch": "querit_search",
|
||||
}
|
||||
for _, name := range names {
|
||||
canonical, ok := canonicalByAlias[name]
|
||||
|
||||
@@ -34,6 +34,9 @@ type ComponentSpec struct {
|
||||
Inputs map[string]string
|
||||
Outputs map[string]string
|
||||
InputForm map[string]any
|
||||
// PreserveJSONNumbers keeps numeric response tokens as json.Number when
|
||||
// Canvas must expose the upstream JSON without float64 precision loss.
|
||||
PreserveJSONNumbers bool
|
||||
}
|
||||
|
||||
// ToolComponent is the required Canvas adaptation contract implemented by a
|
||||
|
||||
@@ -113,6 +113,7 @@ const (
|
||||
EnvSSHEnableAPIURL = "SSH_ENABLE_API_URL"
|
||||
EnvAllowAnyHost = "ALLOW_ANY_HOST"
|
||||
EnvTavilyApiKey = "TAVILY_API_KEY"
|
||||
EnvQueritAPIKey = "QUERIT_API_KEY"
|
||||
EnvHome = "HOME"
|
||||
EnvUserProfile = "USERPROFILE"
|
||||
EnvHttpHTTPProxy = "http_proxy"
|
||||
|
||||
560
test/unit_test/agent/tools/test_querit_unit.py
Normal file
560
test/unit_test/agent/tools/test_querit_unit.py
Normal file
@@ -0,0 +1,560 @@
|
||||
#
|
||||
# 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 pytest
|
||||
|
||||
import agent.tools.querit as querit_module
|
||||
from agent.tools.querit import QueritSearch, QueritSearchParam
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status_code=200):
|
||||
self._payload = payload
|
||||
self.status_code = status_code
|
||||
|
||||
def json(self):
|
||||
return self._payload
|
||||
|
||||
def raise_for_status(self):
|
||||
if self.status_code >= 400:
|
||||
raise querit_module.requests.HTTPError(
|
||||
f"{self.status_code} error",
|
||||
response=self,
|
||||
)
|
||||
|
||||
|
||||
def _make_tool(api_key="test-api-key"):
|
||||
tool = QueritSearch.__new__(QueritSearch)
|
||||
param = QueritSearchParam()
|
||||
param.api_key = api_key
|
||||
param.delay_after_error = 0
|
||||
tool._param = param
|
||||
tool.check_if_canceled = lambda *args, **kwargs: False
|
||||
|
||||
captured = {}
|
||||
outputs = {}
|
||||
|
||||
def fake_retrieve(results, get_title, get_url, get_content, get_score):
|
||||
items = list(results)
|
||||
captured["references"] = [
|
||||
{
|
||||
"title": get_title(item),
|
||||
"url": get_url(item),
|
||||
"content": get_content(item),
|
||||
"score": get_score(item),
|
||||
}
|
||||
for item in items
|
||||
]
|
||||
outputs["formalized_content"] = "FORMALIZED"
|
||||
|
||||
tool._retrieve_chunks = fake_retrieve
|
||||
tool.set_output = lambda key, value: outputs.__setitem__(key, value)
|
||||
tool.output = lambda key=None: outputs.get(key) if key else outputs
|
||||
return tool, captured, outputs
|
||||
|
||||
|
||||
def test_minimal_search_posts_defaults_and_preserves_raw_response(monkeypatch):
|
||||
raw_response = {
|
||||
"took": "300ms",
|
||||
"error_code": 200,
|
||||
"error_msg": "",
|
||||
"search_id": 11099848653006015581,
|
||||
"query_context": {"query": "expanded query", "count": 99},
|
||||
"results": {
|
||||
"result": [
|
||||
{
|
||||
"title": "LLM progress",
|
||||
"url": "https://example.com/llm",
|
||||
"snippet": "Recent LLM progress summary",
|
||||
"custom": {"request_id": "kept"},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
calls = []
|
||||
|
||||
def fake_post(url, **kwargs):
|
||||
calls.append((url, kwargs))
|
||||
return _FakeResponse(raw_response)
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, captured, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="latest LLM progress")
|
||||
|
||||
assert result == "FORMALIZED"
|
||||
assert calls == [
|
||||
(
|
||||
"https://api.querit.ai/v1/search",
|
||||
{
|
||||
"headers": {
|
||||
"Accept": "application/json",
|
||||
"Authorization": "Bearer test-api-key",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"json": {
|
||||
"query": "latest LLM progress",
|
||||
"count": 10,
|
||||
"chunksPerDoc": 3,
|
||||
},
|
||||
"timeout": querit_module.DEFAULT_TIMEOUT,
|
||||
},
|
||||
)
|
||||
]
|
||||
assert captured["references"] == [
|
||||
{
|
||||
"title": "LLM progress",
|
||||
"url": "https://example.com/llm",
|
||||
"content": "Recent LLM progress summary",
|
||||
"score": 1,
|
||||
}
|
||||
]
|
||||
assert outputs["json"] == raw_response
|
||||
assert outputs["json"]["search_id"] == 11099848653006015581
|
||||
|
||||
|
||||
def test_search_maps_flat_filters_to_querit_request(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_post(url, **kwargs):
|
||||
calls.append((url, kwargs))
|
||||
return _FakeResponse({"results": {"result": []}})
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(
|
||||
query="AI search news",
|
||||
count=5,
|
||||
chunks_per_doc=1,
|
||||
site_include=["example.com"],
|
||||
site_exclude=["archive.example.com"],
|
||||
time_range="w1",
|
||||
country_include=["US", "CA"],
|
||||
language_include=["en", "fr"],
|
||||
)
|
||||
|
||||
assert result == ""
|
||||
assert calls[0][1]["json"] == {
|
||||
"query": "AI search news",
|
||||
"count": 5,
|
||||
"chunksPerDoc": 1,
|
||||
"filters": {
|
||||
"sites": {
|
||||
"include": ["example.com"],
|
||||
"exclude": ["archive.example.com"],
|
||||
},
|
||||
"timeRange": {"date": "w1"},
|
||||
"geo": {"countries": {"include": ["US", "CA"]}},
|
||||
"languages": {"include": ["en", "fr"]},
|
||||
},
|
||||
}
|
||||
assert outputs["json"] == {"results": {"result": []}}
|
||||
|
||||
|
||||
def test_unauthorized_response_is_not_retried(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_post(url, **kwargs):
|
||||
calls.append((url, kwargs))
|
||||
return _FakeResponse({"error": "unauthorized"}, status_code=401)
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert len(calls) == 1
|
||||
assert "401" in result
|
||||
assert "401" in outputs["_ERROR"]
|
||||
|
||||
|
||||
def test_blank_node_api_key_falls_back_to_environment(monkeypatch):
|
||||
received_authorization = []
|
||||
|
||||
def fake_post(_url, **kwargs):
|
||||
received_authorization.append(kwargs["headers"]["Authorization"])
|
||||
return _FakeResponse({"results": {"result": []}})
|
||||
|
||||
monkeypatch.setenv("QUERIT_API_KEY", "environment-key")
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool(api_key=" ")
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert result == ""
|
||||
assert outputs.get("_ERROR", "") == ""
|
||||
assert received_authorization == ["Bearer environment-key"]
|
||||
|
||||
|
||||
def test_missing_api_key_fails_before_network(monkeypatch):
|
||||
calls = []
|
||||
monkeypatch.delenv("QUERIT_API_KEY", raising=False)
|
||||
monkeypatch.setattr(querit_module.requests, "post", lambda *args, **kwargs: calls.append((args, kwargs)))
|
||||
tool, _, outputs = _make_tool(api_key="")
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert calls == []
|
||||
assert "QUERIT_API_KEY" in result
|
||||
assert "QUERIT_API_KEY" in outputs["_ERROR"]
|
||||
|
||||
|
||||
def test_empty_query_returns_empty_outputs_without_network(monkeypatch):
|
||||
calls = []
|
||||
monkeypatch.setattr(querit_module.requests, "post", lambda *args, **kwargs: calls.append((args, kwargs)))
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="")
|
||||
|
||||
assert result == ""
|
||||
assert calls == []
|
||||
assert outputs == {"formalized_content": "", "json": {}}
|
||||
|
||||
|
||||
def test_explicit_none_omits_chunks_per_doc(monkeypatch):
|
||||
payloads = []
|
||||
|
||||
def fake_post(_url, **kwargs):
|
||||
payloads.append(kwargs["json"])
|
||||
return _FakeResponse({"results": {"result": []}})
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, _ = _make_tool()
|
||||
|
||||
tool._invoke(query="AI news", chunks_per_doc=None)
|
||||
|
||||
assert payloads == [{"query": "AI news", "count": 10}]
|
||||
|
||||
|
||||
def test_runtime_values_override_node_defaults(monkeypatch):
|
||||
payloads = []
|
||||
|
||||
def fake_post(_url, **kwargs):
|
||||
payloads.append(kwargs["json"])
|
||||
return _FakeResponse({"results": {"result": []}})
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, _ = _make_tool()
|
||||
tool._param.count = 8
|
||||
tool._param.chunks_per_doc = 2
|
||||
tool._param.site_include = ["node.example.com"]
|
||||
|
||||
tool._invoke(
|
||||
query="AI news",
|
||||
count=1,
|
||||
site_include=[],
|
||||
time_range="2026-01-01to2026-01-31",
|
||||
)
|
||||
|
||||
assert payloads == [
|
||||
{
|
||||
"query": "AI news",
|
||||
"count": 1,
|
||||
"chunksPerDoc": 2,
|
||||
"filters": {"timeRange": {"date": "2026-01-01to2026-01-31"}},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("runtime_values", "expected_error"),
|
||||
[
|
||||
({"query": 123}, "query"),
|
||||
({"query": "AI news", "count": 0}, "count"),
|
||||
({"query": "AI news", "count": 1.5}, "count"),
|
||||
({"query": "AI news", "chunks_per_doc": 0}, "chunks_per_doc"),
|
||||
({"query": "AI news", "chunks_per_doc": 4}, "chunks_per_doc"),
|
||||
({"query": "AI news", "time_range": "past_week"}, "time_range"),
|
||||
({"query": "AI news", "site_include": ["example.com", 1]}, "site_include"),
|
||||
],
|
||||
)
|
||||
def test_invalid_runtime_values_fail_before_network(monkeypatch, runtime_values, expected_error):
|
||||
calls = []
|
||||
monkeypatch.setattr(querit_module.requests, "post", lambda *args, **kwargs: calls.append((args, kwargs)))
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(**runtime_values)
|
||||
|
||||
assert calls == []
|
||||
assert expected_error in result
|
||||
assert expected_error in outputs["_ERROR"]
|
||||
|
||||
|
||||
def test_retryable_http_errors_are_retried_until_success(monkeypatch):
|
||||
responses = [
|
||||
_FakeResponse({"error": "rate limited"}, status_code=429),
|
||||
_FakeResponse({"error": "temporary"}, status_code=503),
|
||||
_FakeResponse({"results": {"result": []}, "request_id": "success"}),
|
||||
]
|
||||
calls = []
|
||||
|
||||
def fake_post(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
return responses.pop(0)
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert result == ""
|
||||
assert len(calls) == 3
|
||||
assert outputs["json"]["request_id"] == "success"
|
||||
|
||||
|
||||
def test_retry_wait_uses_configured_delay(monkeypatch):
|
||||
responses = [
|
||||
_FakeResponse({"error": "rate limited"}, status_code=429),
|
||||
_FakeResponse({"results": {"result": []}}),
|
||||
]
|
||||
delays = []
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", lambda *args, **kwargs: responses.pop(0))
|
||||
monkeypatch.setattr(querit_module.time, "sleep", delays.append)
|
||||
tool, _, _ = _make_tool()
|
||||
tool._param.delay_after_error = 0.25
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert result == ""
|
||||
assert delays == [0.25]
|
||||
|
||||
|
||||
def test_cancellation_stops_before_retry(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_post(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
return _FakeResponse({"error": "rate limited"}, status_code=429)
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool()
|
||||
cancellation_checks = iter([False, False, True])
|
||||
tool.check_if_canceled = lambda *args, **kwargs: next(cancellation_checks)
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert result is None
|
||||
assert len(calls) == 1
|
||||
assert "_ERROR" not in outputs
|
||||
|
||||
|
||||
def test_network_errors_are_retried_until_success(monkeypatch):
|
||||
outcomes = [
|
||||
querit_module.requests.ConnectTimeout("first timeout"),
|
||||
querit_module.requests.ConnectionError("second failure"),
|
||||
_FakeResponse({"results": {"result": []}}),
|
||||
]
|
||||
calls = []
|
||||
|
||||
def fake_post(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
outcome = outcomes.pop(0)
|
||||
if isinstance(outcome, Exception):
|
||||
raise outcome
|
||||
return outcome
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert result == ""
|
||||
assert len(calls) == 3
|
||||
assert outputs["json"] == {"results": {"result": []}}
|
||||
|
||||
|
||||
def test_invalid_json_response_is_not_retried(monkeypatch):
|
||||
calls = []
|
||||
|
||||
class _InvalidJSONResponse(_FakeResponse):
|
||||
def json(self):
|
||||
raise querit_module.requests.JSONDecodeError("invalid JSON", "not-json", 0)
|
||||
|
||||
def fake_post(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
return _InvalidJSONResponse(None)
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert len(calls) == 1
|
||||
assert "invalid JSON" in result
|
||||
assert "invalid JSON" in outputs["_ERROR"]
|
||||
|
||||
|
||||
def test_persistent_retryable_error_stops_after_three_attempts(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_post(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
return _FakeResponse({"error": "temporary"}, status_code=500)
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert len(calls) == 3
|
||||
assert "500" in result
|
||||
assert "500" in outputs["_ERROR"]
|
||||
|
||||
|
||||
def test_non_object_response_is_rejected(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
querit_module.requests,
|
||||
"post",
|
||||
lambda *args, **kwargs: _FakeResponse(["unexpected"]),
|
||||
)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert "JSON object" in result
|
||||
assert "JSON object" in outputs["_ERROR"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw_response", "expected_error"),
|
||||
[
|
||||
({"results": []}, "results must be an object"),
|
||||
({"results": None}, "results must be an object"),
|
||||
({"results": {"result": {}}}, "results.result must be an array"),
|
||||
({"results": {"result": None}}, "results.result must be an array"),
|
||||
],
|
||||
)
|
||||
def test_invalid_result_container_types_are_rejected(
|
||||
monkeypatch,
|
||||
raw_response,
|
||||
expected_error,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
querit_module.requests,
|
||||
"post",
|
||||
lambda *args, **kwargs: _FakeResponse(raw_response),
|
||||
)
|
||||
tool, _, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert expected_error in result
|
||||
assert expected_error in outputs["_ERROR"]
|
||||
|
||||
|
||||
def test_results_with_missing_optional_fields_do_not_fail(monkeypatch):
|
||||
raw_response = {
|
||||
"results": {
|
||||
"result": [
|
||||
{"title": "Title only"},
|
||||
{"url": "https://example.com"},
|
||||
{"snippet": "Snippet only"},
|
||||
"unexpected",
|
||||
]
|
||||
}
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
querit_module.requests,
|
||||
"post",
|
||||
lambda *args, **kwargs: _FakeResponse(raw_response),
|
||||
)
|
||||
tool, captured, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert result == "FORMALIZED"
|
||||
assert len(captured["references"]) == 3
|
||||
assert outputs["json"] == raw_response
|
||||
|
||||
|
||||
def test_non_string_reference_fields_match_go_string_coercion(monkeypatch):
|
||||
raw_response = {
|
||||
"results": {
|
||||
"result": [
|
||||
{
|
||||
"title": 123,
|
||||
"url": 456,
|
||||
"snippet": 789,
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
querit_module.requests,
|
||||
"post",
|
||||
lambda *args, **kwargs: _FakeResponse(raw_response),
|
||||
)
|
||||
tool, captured, outputs = _make_tool()
|
||||
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert result == "FORMALIZED"
|
||||
assert captured["references"] == [
|
||||
{
|
||||
"title": "123",
|
||||
"url": "456",
|
||||
"content": "789",
|
||||
"score": 1,
|
||||
}
|
||||
]
|
||||
assert outputs["json"] == raw_response
|
||||
|
||||
|
||||
def test_api_key_is_redacted_from_errors_and_logs(monkeypatch, caplog):
|
||||
secret = f"secret-test-key-{id(caplog)}"
|
||||
|
||||
def fake_post(*args, **kwargs):
|
||||
raise querit_module.requests.ConnectionError(f"connection failed for {secret}")
|
||||
|
||||
monkeypatch.setattr(querit_module.requests, "post", fake_post)
|
||||
tool, _, outputs = _make_tool(api_key=secret)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
result = tool._invoke(query="AI news")
|
||||
|
||||
assert secret not in result
|
||||
assert secret not in outputs["_ERROR"]
|
||||
assert secret not in caplog.text
|
||||
assert "[REDACTED]" in result
|
||||
|
||||
|
||||
def test_querit_classes_are_available_through_dynamic_tool_discovery():
|
||||
import agent.tools as tools_package
|
||||
|
||||
assert tools_package.QueritSearch is QueritSearch
|
||||
assert tools_package.QueritSearchParam is QueritSearchParam
|
||||
|
||||
|
||||
def test_tool_metadata_exposes_only_supported_runtime_parameters():
|
||||
metadata = QueritSearchParam().get_meta()["function"]
|
||||
|
||||
assert metadata["name"] == "querit_search"
|
||||
assert metadata["parameters"]["required"] == ["query"]
|
||||
assert set(metadata["parameters"]["properties"]) == {
|
||||
"query",
|
||||
"count",
|
||||
"chunks_per_doc",
|
||||
"site_include",
|
||||
"site_exclude",
|
||||
"time_range",
|
||||
"country_include",
|
||||
"language_include",
|
||||
}
|
||||
assert "api_key" not in metadata["parameters"]["properties"]
|
||||
@@ -159,7 +159,7 @@ func sortedToolNames() []string {
|
||||
known := []string{
|
||||
"akshare", "arxiv", "code_exec", "crawler", "deepl", "duckduckgo",
|
||||
"email", "execute_sql", "exesql", "github", "google", "google_scholar",
|
||||
"jin10", "pubmed", "qweather", "retrieval", "search_my_dataset",
|
||||
"jin10", "pubmed", "qweather", "querit", "querit_search", "queritsearch", "retrieval", "search_my_dataset",
|
||||
"search_my_dateset", "searxng", "tavily", "tushare", "wencai",
|
||||
"web_crawler", "wikipedia", "yahoo_finance",
|
||||
}
|
||||
|
||||
BIN
web/src/assets/querit.png
Normal file
BIN
web/src/assets/querit.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 48 KiB |
@@ -95,6 +95,7 @@ export enum Operator {
|
||||
Tool = 'Tool',
|
||||
TavilySearch = 'TavilySearch',
|
||||
TavilyExtract = 'TavilyExtract',
|
||||
QueritSearch = 'QueritSearch',
|
||||
UserFillUp = 'UserFillUp',
|
||||
StringTransform = 'StringTransform',
|
||||
SearXNG = 'SearXNG',
|
||||
|
||||
@@ -2491,6 +2491,26 @@ Best for: Documents with flowing, contextually connected content — such as boo
|
||||
includeImageDescriptions: 'Include image descriptions',
|
||||
includeDomains: 'Include domains',
|
||||
ExcludeDomains: 'Exclude domains',
|
||||
queritCount: 'Result count',
|
||||
queritCountTip: 'Set the maximum number of search results to return.',
|
||||
queritChunksPerDoc: 'Chunks per document',
|
||||
queritChunksPerDocTip:
|
||||
'Set the number of relevant text chunks returned for each result, from 1 to 3.',
|
||||
queritSiteInclude: 'Sites to include',
|
||||
queritSiteIncludeTip:
|
||||
'Return results only from the specified domains or websites.',
|
||||
queritSiteExclude: 'Sites to exclude',
|
||||
queritSiteExcludeTip:
|
||||
'Exclude results from the specified domains or websites.',
|
||||
queritTimeRange: 'Time range',
|
||||
queritTimeRangeTip:
|
||||
'Use a relative range such as d7, w2, m1, or y1, or an absolute range such as 2026-01-01to2026-01-31.',
|
||||
queritCountryInclude: 'Countries to include',
|
||||
queritCountryIncludeTip:
|
||||
'Return results associated with the specified countries.',
|
||||
queritLanguageInclude: 'Languages to include',
|
||||
queritLanguageIncludeTip: 'Return results in the specified languages.',
|
||||
queritListPlaceholder: 'Enter a value',
|
||||
Days: 'Days',
|
||||
comma: 'Comma',
|
||||
semicolon: 'Semicolon',
|
||||
@@ -3147,6 +3167,9 @@ This process aggregates variables from multiple branches into a single variable
|
||||
codeExec: 'Code',
|
||||
tavilySearch: 'Tavily search',
|
||||
tavilySearchDescription: 'Search results via Tavily service.',
|
||||
queritSearch: 'Querit search',
|
||||
queritSearchDescription:
|
||||
'Search the web with Querit and return source-backed results for agents.',
|
||||
tavilyExtract: 'Tavily extract',
|
||||
tavilyExtractDescription: 'Tavily Extract',
|
||||
log: 'Log',
|
||||
@@ -3179,6 +3202,7 @@ This process aggregates variables from multiple branches into a single variable
|
||||
code: 'Running a quick script',
|
||||
textProcessing: 'Tidying up text',
|
||||
tavilySearch: 'Searching the web',
|
||||
queritSearch: 'Searching the web with Querit',
|
||||
tavilyExtract: 'Reading the page',
|
||||
exeSQL: 'Querying database',
|
||||
google: 'Searching the web',
|
||||
|
||||
@@ -2141,6 +2141,23 @@ NER:使用 spaCy NER 和基于规则的关键词提取来抽取实体和关系
|
||||
includeImageDescriptions: '包含图片描述',
|
||||
includeDomains: '包含域名',
|
||||
ExcludeDomains: '排除域名',
|
||||
queritCount: '结果数量',
|
||||
queritCountTip: '设置本次搜索最多返回的结果数量。',
|
||||
queritChunksPerDoc: '每篇文档片段数',
|
||||
queritChunksPerDocTip:
|
||||
'设置每条结果返回的相关文本片段数量,可填写 1 至 3。',
|
||||
queritSiteInclude: '包含的网站',
|
||||
queritSiteIncludeTip: '仅返回来自指定域名或网站的结果。',
|
||||
queritSiteExclude: '排除的网站',
|
||||
queritSiteExcludeTip: '排除来自指定域名或网站的结果。',
|
||||
queritTimeRange: '时间范围',
|
||||
queritTimeRangeTip:
|
||||
'可填写 d7、w2、m1、y1 等相对时间,或 2026-01-01to2026-01-31 形式的绝对日期范围。',
|
||||
queritCountryInclude: '包含的国家',
|
||||
queritCountryIncludeTip: '返回与指定国家相关的结果。',
|
||||
queritLanguageInclude: '包含的语言',
|
||||
queritLanguageIncludeTip: '返回使用指定语言的结果。',
|
||||
queritListPlaceholder: '请输入内容',
|
||||
days: '天数',
|
||||
comma: '逗号',
|
||||
semicolon: '分号',
|
||||
@@ -2758,6 +2775,9 @@ NER:使用 spaCy NER 和基于规则的关键词提取来抽取实体和关系
|
||||
codeExec: '代码',
|
||||
tavilySearch: 'Tavily 搜索',
|
||||
tavilySearchDescription: '通过 Tavily 服务搜索结果',
|
||||
queritSearch: 'Querit 搜索',
|
||||
queritSearchDescription:
|
||||
'使用 Querit 搜索网络,并为智能体返回带来源的搜索结果。',
|
||||
tavilyExtract: 'Tavily 提取',
|
||||
tavilyExtractDescription: 'Tavily 提取',
|
||||
log: '日志',
|
||||
@@ -2778,6 +2798,7 @@ NER:使用 spaCy NER 和基于规则的关键词提取来抽取实体和关系
|
||||
code: '运行小段代码',
|
||||
textProcessing: '整理文字',
|
||||
tavilySearch: '正在网上搜索',
|
||||
queritSearch: '正在使用 Querit 搜索网络',
|
||||
tavilyExtract: '读取网页内容',
|
||||
exeSQL: '查询数据库',
|
||||
google: '正在网上搜索',
|
||||
|
||||
@@ -109,6 +109,7 @@ export function AccordionOperators({
|
||||
operators={[
|
||||
Operator.TavilySearch,
|
||||
Operator.TavilyExtract,
|
||||
Operator.QueritSearch,
|
||||
Operator.ExeSQL,
|
||||
Operator.Google,
|
||||
Operator.YahooFinance,
|
||||
|
||||
@@ -572,6 +572,28 @@ export const initialTavilyValues = {
|
||||
},
|
||||
};
|
||||
|
||||
export const initialQueritValues = {
|
||||
api_key: '',
|
||||
query: AgentGlobals.SysQuery,
|
||||
count: 10,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
outputs: {
|
||||
formalized_content: {
|
||||
value: '',
|
||||
type: 'string',
|
||||
},
|
||||
json: {
|
||||
value: {},
|
||||
type: 'object',
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
export enum TavilyExtractDepth {
|
||||
Basic = 'basic',
|
||||
Advanced = 'advanced',
|
||||
@@ -722,6 +744,7 @@ export const RestrictedUpstreamMap = {
|
||||
[Operator.WaitingDialogue]: [Operator.Begin],
|
||||
[Operator.Agent]: [Operator.Begin],
|
||||
[Operator.TavilySearch]: [Operator.Begin],
|
||||
[Operator.QueritSearch]: [Operator.Begin],
|
||||
[Operator.TavilyExtract]: [Operator.Begin],
|
||||
[Operator.StringTransform]: [Operator.Begin],
|
||||
[Operator.UserFillUp]: [Operator.Begin],
|
||||
@@ -777,6 +800,7 @@ export const NodeMap = {
|
||||
[Operator.Agent]: 'agentNode',
|
||||
[Operator.Tool]: 'toolNode',
|
||||
[Operator.TavilySearch]: 'ragNode',
|
||||
[Operator.QueritSearch]: 'ragNode',
|
||||
[Operator.UserFillUp]: 'ragNode',
|
||||
[Operator.StringTransform]: 'ragNode',
|
||||
[Operator.TavilyExtract]: 'ragNode',
|
||||
|
||||
@@ -26,6 +26,7 @@ import LoopForm from '../form/loop-form';
|
||||
import MessageForm from '../form/message-form';
|
||||
import ParserForm from '../form/parser-form';
|
||||
import PubMedForm from '../form/pubmed-form';
|
||||
import QueritForm from '../form/querit-form';
|
||||
import BGPTForm from '../form/bgpt-form';
|
||||
import RetrievalForm from '../form/retrieval-form/next';
|
||||
import RewriteQuestionForm from '../form/rewrite-question-form';
|
||||
@@ -145,6 +146,9 @@ export const FormConfigMap = {
|
||||
[Operator.TavilySearch]: {
|
||||
component: TavilyForm,
|
||||
},
|
||||
[Operator.QueritSearch]: {
|
||||
component: QueritForm,
|
||||
},
|
||||
[Operator.UserFillUp]: {
|
||||
component: UserFillUpForm,
|
||||
},
|
||||
|
||||
@@ -23,6 +23,7 @@ const Menus = [
|
||||
list: [
|
||||
Operator.TavilySearch,
|
||||
Operator.TavilyExtract,
|
||||
Operator.QueritSearch,
|
||||
Operator.Google,
|
||||
// Operator.Bing,
|
||||
Operator.DuckDuckGo,
|
||||
|
||||
73
web/src/pages/agent/form/querit-form/dynamic-list.tsx
Normal file
73
web/src/pages/agent/form/querit-form/dynamic-list.tsx
Normal file
@@ -0,0 +1,73 @@
|
||||
import { BlockButton, Button } from '@/components/ui/button';
|
||||
import {
|
||||
FormControl,
|
||||
FormField,
|
||||
FormItem,
|
||||
FormLabel,
|
||||
FormMessage,
|
||||
} from '@/components/ui/form';
|
||||
import { Input } from '@/components/ui/input';
|
||||
import { X } from 'lucide-react';
|
||||
import { useFieldArray, useFormContext } from 'react-hook-form';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
type DynamicListProps = {
|
||||
name:
|
||||
| 'site_include'
|
||||
| 'site_exclude'
|
||||
| 'country_include'
|
||||
| 'language_include';
|
||||
label: string;
|
||||
tooltip: string;
|
||||
placeholder: string;
|
||||
};
|
||||
|
||||
export function DynamicList({
|
||||
name,
|
||||
label,
|
||||
tooltip,
|
||||
placeholder,
|
||||
}: DynamicListProps) {
|
||||
const { t } = useTranslation();
|
||||
const form = useFormContext();
|
||||
const { fields, append, remove } = useFieldArray({
|
||||
name,
|
||||
control: form.control,
|
||||
});
|
||||
|
||||
return (
|
||||
<FormItem>
|
||||
<FormLabel tooltip={tooltip}>{label}</FormLabel>
|
||||
<div className="space-y-4">
|
||||
{fields.map((item, index) => (
|
||||
<div key={item.id} className="flex items-start gap-2">
|
||||
<FormField
|
||||
control={form.control}
|
||||
name={`${name}.${index}.value`}
|
||||
render={({ field }) => (
|
||||
<FormItem className="flex-1">
|
||||
<FormControl>
|
||||
<Input {...field} placeholder={placeholder} />
|
||||
</FormControl>
|
||||
<FormMessage />
|
||||
</FormItem>
|
||||
)}
|
||||
/>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
aria-label={`${t('common.remove')} ${label} ${index + 1}`}
|
||||
onClick={() => remove(index)}
|
||||
>
|
||||
<X />
|
||||
</Button>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
<BlockButton type="button" onClick={() => append({ value: '' })}>
|
||||
{t('common.add')}
|
||||
</BlockButton>
|
||||
</FormItem>
|
||||
);
|
||||
}
|
||||
136
web/src/pages/agent/form/querit-form/index.tsx
Normal file
136
web/src/pages/agent/form/querit-form/index.tsx
Normal file
@@ -0,0 +1,136 @@
|
||||
import { FormContainer } from '@/components/form-container';
|
||||
import {
|
||||
Form,
|
||||
FormControl,
|
||||
FormField,
|
||||
FormItem,
|
||||
FormLabel,
|
||||
FormMessage,
|
||||
} from '@/components/ui/form';
|
||||
import { Input } from '@/components/ui/input';
|
||||
import { useTranslate } from '@/hooks/common-hooks';
|
||||
import { zodResolver } from '@hookform/resolvers/zod';
|
||||
import { memo } from 'react';
|
||||
import { useForm, useFormContext } from 'react-hook-form';
|
||||
import { z } from 'zod';
|
||||
import { initialQueritValues } from '../../constant';
|
||||
import { INextOperatorForm } from '../../interface';
|
||||
import { buildOutputList } from '../../utils/build-output-list';
|
||||
import { ApiKeyField } from '../components/api-key-field';
|
||||
import { FormWrapper } from '../components/form-wrapper';
|
||||
import { Output } from '../components/output';
|
||||
import { QueryVariable } from '../components/query-variable';
|
||||
import { DynamicList } from './dynamic-list';
|
||||
import { QueritFormSchema } from './utils';
|
||||
import { useValues } from './use-values';
|
||||
import { useWatchFormChange } from './use-watch-change';
|
||||
|
||||
const outputList = buildOutputList(initialQueritValues.outputs);
|
||||
|
||||
export function QueritFormWidgets({ includeQuery = true }) {
|
||||
const form = useFormContext();
|
||||
const { t } = useTranslate('flow');
|
||||
const listPlaceholder = t('queritListPlaceholder');
|
||||
|
||||
return (
|
||||
<>
|
||||
<ApiKeyField />
|
||||
{includeQuery && <QueryVariable />}
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="count"
|
||||
render={({ field }) => (
|
||||
<FormItem>
|
||||
<FormLabel tooltip={t('queritCountTip')}>
|
||||
{t('queritCount')}
|
||||
</FormLabel>
|
||||
<FormControl>
|
||||
<Input type="number" min={1} step={1} {...field} />
|
||||
</FormControl>
|
||||
<FormMessage />
|
||||
</FormItem>
|
||||
)}
|
||||
/>
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="chunks_per_doc"
|
||||
render={({ field }) => (
|
||||
<FormItem>
|
||||
<FormLabel tooltip={t('queritChunksPerDocTip')}>
|
||||
{t('queritChunksPerDoc')}
|
||||
</FormLabel>
|
||||
<FormControl>
|
||||
<Input type="number" min={1} max={3} step={1} {...field} />
|
||||
</FormControl>
|
||||
<FormMessage />
|
||||
</FormItem>
|
||||
)}
|
||||
/>
|
||||
<DynamicList
|
||||
name="site_include"
|
||||
label={t('queritSiteInclude')}
|
||||
tooltip={t('queritSiteIncludeTip')}
|
||||
placeholder={listPlaceholder}
|
||||
/>
|
||||
<DynamicList
|
||||
name="site_exclude"
|
||||
label={t('queritSiteExclude')}
|
||||
tooltip={t('queritSiteExcludeTip')}
|
||||
placeholder={listPlaceholder}
|
||||
/>
|
||||
<FormField
|
||||
control={form.control}
|
||||
name="time_range"
|
||||
render={({ field }) => (
|
||||
<FormItem>
|
||||
<FormLabel tooltip={t('queritTimeRangeTip')}>
|
||||
{t('queritTimeRange')}
|
||||
</FormLabel>
|
||||
<FormControl>
|
||||
<Input placeholder="d7" {...field} />
|
||||
</FormControl>
|
||||
<FormMessage />
|
||||
</FormItem>
|
||||
)}
|
||||
/>
|
||||
<DynamicList
|
||||
name="country_include"
|
||||
label={t('queritCountryInclude')}
|
||||
tooltip={t('queritCountryIncludeTip')}
|
||||
placeholder={listPlaceholder}
|
||||
/>
|
||||
<DynamicList
|
||||
name="language_include"
|
||||
label={t('queritLanguageInclude')}
|
||||
tooltip={t('queritLanguageIncludeTip')}
|
||||
placeholder={listPlaceholder}
|
||||
/>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
function QueritForm({ node }: INextOperatorForm) {
|
||||
const values = useValues(node);
|
||||
const form = useForm<z.infer<typeof QueritFormSchema>>({
|
||||
defaultValues: values,
|
||||
resolver: zodResolver(QueritFormSchema),
|
||||
mode: 'onChange',
|
||||
});
|
||||
|
||||
useWatchFormChange(node?.id, form);
|
||||
|
||||
return (
|
||||
<Form {...form}>
|
||||
<FormWrapper>
|
||||
<FormContainer>
|
||||
<QueritFormWidgets />
|
||||
</FormContainer>
|
||||
</FormWrapper>
|
||||
<div className="p-5">
|
||||
<Output list={outputList} />
|
||||
</div>
|
||||
</Form>
|
||||
);
|
||||
}
|
||||
|
||||
export default memo(QueritForm);
|
||||
14
web/src/pages/agent/form/querit-form/use-values.ts
Normal file
14
web/src/pages/agent/form/querit-form/use-values.ts
Normal file
@@ -0,0 +1,14 @@
|
||||
import { RAGFlowNodeType } from '@/interfaces/database/agent';
|
||||
import { isEmpty } from 'lodash';
|
||||
import { useMemo } from 'react';
|
||||
import { initialQueritValues } from '../../constant';
|
||||
import { deserializeQueritFormValues } from './utils';
|
||||
|
||||
export function useValues(node?: RAGFlowNodeType) {
|
||||
return useMemo(() => {
|
||||
const formData = node?.data?.form;
|
||||
const values = isEmpty(formData) ? initialQueritValues : formData;
|
||||
|
||||
return deserializeQueritFormValues(values);
|
||||
}, [node?.data?.form]);
|
||||
}
|
||||
114
web/src/pages/agent/form/querit-form/use-watch-change.test.tsx
Normal file
114
web/src/pages/agent/form/querit-form/use-watch-change.test.tsx
Normal file
@@ -0,0 +1,114 @@
|
||||
import { act, renderHook, waitFor } from '@testing-library/react';
|
||||
import { StrictMode } from 'react';
|
||||
import { useForm } from 'react-hook-form';
|
||||
import { useWatchFormChange } from './use-watch-change';
|
||||
|
||||
const mockUpdateNodeForm = jest.fn();
|
||||
|
||||
jest.mock('../../store', () => ({
|
||||
__esModule: true,
|
||||
default: (selector: (state: Record<string, unknown>) => unknown) =>
|
||||
selector({ updateNodeForm: mockUpdateNodeForm }),
|
||||
}));
|
||||
|
||||
const validValues = {
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 10,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
};
|
||||
|
||||
function useQueritFormWatcher(defaultValues: Record<string, unknown>) {
|
||||
const form = useForm({ defaultValues });
|
||||
useWatchFormChange('querit:0', form);
|
||||
return form;
|
||||
}
|
||||
|
||||
describe('useWatchFormChange', () => {
|
||||
beforeEach(() => {
|
||||
mockUpdateNodeForm.mockClear();
|
||||
});
|
||||
|
||||
it('skips the initial write and normalizes subsequent valid edits', async () => {
|
||||
const { result } = renderHook(() =>
|
||||
useQueritFormWatcher({
|
||||
...validValues,
|
||||
count: '5',
|
||||
chunks_per_doc: '2',
|
||||
site_include: [{ value: ' docs.example.com ' }],
|
||||
}),
|
||||
);
|
||||
|
||||
expect(mockUpdateNodeForm).not.toHaveBeenCalled();
|
||||
|
||||
act(() => {
|
||||
result.current.setValue('count', '6');
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUpdateNodeForm).toHaveBeenCalledWith(
|
||||
'querit:0',
|
||||
expect.objectContaining({
|
||||
count: 6,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it('does not write an invalid initial form to the Canvas node', async () => {
|
||||
renderHook(() =>
|
||||
useQueritFormWatcher({
|
||||
...validValues,
|
||||
count: 0,
|
||||
}),
|
||||
);
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(mockUpdateNodeForm).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('does not write a valid initial form in Strict Mode', async () => {
|
||||
renderHook(() => useQueritFormWatcher(validValues), {
|
||||
wrapper: StrictMode,
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(mockUpdateNodeForm).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('does not overwrite the last valid state after an invalid edit', async () => {
|
||||
const { result } = renderHook(() => useQueritFormWatcher(validValues));
|
||||
|
||||
act(() => {
|
||||
result.current.setValue('count', 5);
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUpdateNodeForm).toHaveBeenCalled();
|
||||
});
|
||||
mockUpdateNodeForm.mockClear();
|
||||
|
||||
act(() => {
|
||||
result.current.setValue('time_range', 'last week');
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(mockUpdateNodeForm).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
25
web/src/pages/agent/form/querit-form/use-watch-change.ts
Normal file
25
web/src/pages/agent/form/querit-form/use-watch-change.ts
Normal file
@@ -0,0 +1,25 @@
|
||||
import { useEffect, useRef } from 'react';
|
||||
import { UseFormReturn, useWatch } from 'react-hook-form';
|
||||
import useGraphStore from '../../store';
|
||||
import { validateQueritFormValuesForPersistence } from './utils';
|
||||
|
||||
export function useWatchFormChange(id?: string, form?: UseFormReturn<any>) {
|
||||
const values = useWatch({ control: form?.control });
|
||||
const updateNodeForm = useGraphStore((state) => state.updateNodeForm);
|
||||
const previousValues = useRef(values);
|
||||
|
||||
useEffect(() => {
|
||||
if (Object.is(previousValues.current, values)) {
|
||||
return;
|
||||
}
|
||||
previousValues.current = values;
|
||||
if (id && form) {
|
||||
const nextValues = validateQueritFormValuesForPersistence(
|
||||
form.getValues(),
|
||||
);
|
||||
if (nextValues) {
|
||||
updateNodeForm(id, nextValues);
|
||||
}
|
||||
}
|
||||
}, [form, id, updateNodeForm, values]);
|
||||
}
|
||||
224
web/src/pages/agent/form/querit-form/utils.test.ts
Normal file
224
web/src/pages/agent/form/querit-form/utils.test.ts
Normal file
@@ -0,0 +1,224 @@
|
||||
import {
|
||||
QueritFormSchema,
|
||||
deserializeQueritFormValues,
|
||||
serializeQueritFormValues,
|
||||
validateQueritFormValuesForPersistence,
|
||||
} from './utils';
|
||||
|
||||
describe('Querit form contract', () => {
|
||||
const validValues = {
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 10,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
};
|
||||
|
||||
it.each(['d7', 'w2', 'm1', 'y1', '2026-01-01to2026-01-31', ''])(
|
||||
'accepts the supported time range %p',
|
||||
(timeRange) => {
|
||||
expect(
|
||||
QueritFormSchema.safeParse({
|
||||
...validValues,
|
||||
time_range: timeRange,
|
||||
}).success,
|
||||
).toBe(true);
|
||||
},
|
||||
);
|
||||
|
||||
it.each(['d0', '7d', '2026-01-01-2026-01-31', 'last week'])(
|
||||
'rejects the unsupported time range %p',
|
||||
(timeRange) => {
|
||||
expect(
|
||||
QueritFormSchema.safeParse({
|
||||
...validValues,
|
||||
time_range: timeRange,
|
||||
}).success,
|
||||
).toBe(false);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([0, -1, 1.5])('rejects the invalid result count %p', (count) => {
|
||||
expect(
|
||||
QueritFormSchema.safeParse({
|
||||
...validValues,
|
||||
count,
|
||||
}).success,
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it.each([0, 1.5, 4])(
|
||||
'rejects the invalid chunks per document value %p',
|
||||
(chunksPerDoc) => {
|
||||
expect(
|
||||
QueritFormSchema.safeParse({
|
||||
...validValues,
|
||||
chunks_per_doc: chunksPerDoc,
|
||||
}).success,
|
||||
).toBe(false);
|
||||
},
|
||||
);
|
||||
|
||||
it.each(['', ' '])('rejects an empty dynamic list item %p', (listItem) => {
|
||||
expect(
|
||||
QueritFormSchema.safeParse({
|
||||
...validValues,
|
||||
site_include: [{ value: listItem }],
|
||||
}).success,
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it.each(['', ' '])('rejects an empty query %p', (query) => {
|
||||
expect(
|
||||
QueritFormSchema.safeParse({
|
||||
...validValues,
|
||||
query,
|
||||
}).success,
|
||||
).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Querit form persistence', () => {
|
||||
it('rehydrates persisted string arrays for dynamic form fields', () => {
|
||||
expect(
|
||||
deserializeQueritFormValues({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
site_exclude: ['spam.example.com'],
|
||||
time_range: 'w1',
|
||||
country_include: ['united states'],
|
||||
language_include: ['en'],
|
||||
}),
|
||||
).toMatchObject({
|
||||
site_include: [{ value: 'docs.example.com' }],
|
||||
site_exclude: [{ value: 'spam.example.com' }],
|
||||
country_include: [{ value: 'united states' }],
|
||||
language_include: [{ value: 'en' }],
|
||||
});
|
||||
});
|
||||
|
||||
it('serializes dynamic form fields as string arrays', () => {
|
||||
expect(
|
||||
serializeQueritFormValues({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: [{ value: 'docs.example.com' }],
|
||||
site_exclude: [{ value: 'spam.example.com' }],
|
||||
time_range: 'w1',
|
||||
country_include: [{ value: 'united states' }],
|
||||
language_include: [{ value: 'en' }],
|
||||
}),
|
||||
).toEqual({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
site_exclude: ['spam.example.com'],
|
||||
time_range: 'w1',
|
||||
country_include: ['united states'],
|
||||
language_include: ['en'],
|
||||
});
|
||||
});
|
||||
|
||||
it('serializes valid numeric input values as numbers', () => {
|
||||
expect(
|
||||
serializeQueritFormValues({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: '5',
|
||||
chunks_per_doc: '2',
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
}),
|
||||
).toMatchObject({
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
});
|
||||
});
|
||||
|
||||
it('does not produce persisted values when the current form is invalid', () => {
|
||||
expect(
|
||||
validateQueritFormValuesForPersistence({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 0,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
}),
|
||||
).toBeUndefined();
|
||||
});
|
||||
|
||||
it('produces normalized persisted values when the current form is valid', () => {
|
||||
expect(
|
||||
validateQueritFormValuesForPersistence({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: '5',
|
||||
chunks_per_doc: '2',
|
||||
site_include: [{ value: 'docs.example.com' }],
|
||||
site_exclude: [],
|
||||
time_range: 'w1',
|
||||
country_include: [],
|
||||
language_include: [{ value: 'en' }],
|
||||
}),
|
||||
).toMatchObject({
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
language_include: ['en'],
|
||||
});
|
||||
});
|
||||
|
||||
it('trims dynamic list values before persistence', () => {
|
||||
expect(
|
||||
validateQueritFormValuesForPersistence({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: [{ value: ' docs.example.com ' }],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [{ value: ' en ' }],
|
||||
}),
|
||||
).toMatchObject({
|
||||
site_include: ['docs.example.com'],
|
||||
language_include: ['en'],
|
||||
});
|
||||
});
|
||||
|
||||
it('normalizes missing persisted list fields to empty arrays', () => {
|
||||
expect(
|
||||
deserializeQueritFormValues({
|
||||
api_key: '',
|
||||
query: 'begin.query',
|
||||
count: 10,
|
||||
chunks_per_doc: 3,
|
||||
time_range: '',
|
||||
}),
|
||||
).toMatchObject({
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
});
|
||||
});
|
||||
});
|
||||
72
web/src/pages/agent/form/querit-form/utils.ts
Normal file
72
web/src/pages/agent/form/querit-form/utils.ts
Normal file
@@ -0,0 +1,72 @@
|
||||
import { z } from 'zod';
|
||||
|
||||
export const QueritTimeRangePattern =
|
||||
/^([dwmy][1-9][0-9]*|\d{4}-\d{2}-\d{2}to\d{4}-\d{2}-\d{2})$/;
|
||||
|
||||
const dynamicStringListSchema = z.array(
|
||||
z.object({
|
||||
value: z.string().trim().min(1),
|
||||
}),
|
||||
);
|
||||
|
||||
export const QueritFormSchema = z.object({
|
||||
api_key: z.string(),
|
||||
query: z.string().trim().min(1),
|
||||
count: z.coerce.number().int().min(1),
|
||||
chunks_per_doc: z.coerce.number().int().min(1).max(3),
|
||||
site_include: dynamicStringListSchema,
|
||||
site_exclude: dynamicStringListSchema,
|
||||
time_range: z
|
||||
.string()
|
||||
.refine((value) => value === '' || QueritTimeRangePattern.test(value)),
|
||||
country_include: dynamicStringListSchema,
|
||||
language_include: dynamicStringListSchema,
|
||||
});
|
||||
|
||||
type QueritFormRecord = Record<string, any>;
|
||||
|
||||
function toObjectArray(values: unknown) {
|
||||
return Array.isArray(values) ? values.map((value) => ({ value })) : [];
|
||||
}
|
||||
|
||||
function toStringArray(values: unknown) {
|
||||
return Array.isArray(values)
|
||||
? values.map((item) => item?.value as string)
|
||||
: [];
|
||||
}
|
||||
|
||||
export function deserializeQueritFormValues(values: QueritFormRecord) {
|
||||
return {
|
||||
...values,
|
||||
site_include: toObjectArray(values.site_include),
|
||||
site_exclude: toObjectArray(values.site_exclude),
|
||||
country_include: toObjectArray(values.country_include),
|
||||
language_include: toObjectArray(values.language_include),
|
||||
};
|
||||
}
|
||||
|
||||
export function serializeQueritFormValues(values: QueritFormRecord) {
|
||||
return {
|
||||
...values,
|
||||
count: Number(values.count),
|
||||
chunks_per_doc: Number(values.chunks_per_doc),
|
||||
site_include: toStringArray(values.site_include),
|
||||
site_exclude: toStringArray(values.site_exclude),
|
||||
country_include: toStringArray(values.country_include),
|
||||
language_include: toStringArray(values.language_include),
|
||||
};
|
||||
}
|
||||
|
||||
export function validateQueritFormValuesForPersistence(
|
||||
values: QueritFormRecord,
|
||||
) {
|
||||
const result = QueritFormSchema.safeParse(values);
|
||||
if (!result.success) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return serializeQueritFormValues({
|
||||
...values,
|
||||
...result.data,
|
||||
});
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import GoogleForm from './google-form';
|
||||
import GoogleScholarForm from './google-scholar-form';
|
||||
import KeenableForm from './keenable-form';
|
||||
import PubMedForm from './pubmed-form';
|
||||
import QueritForm from './querit-form';
|
||||
import BGPTForm from './bgpt-form';
|
||||
import RetrievalForm from './retrieval-form';
|
||||
import SearXNGForm from './searxng-form';
|
||||
@@ -36,6 +37,7 @@ export const ToolFormConfigMap = {
|
||||
[Operator.Email]: EmailForm,
|
||||
[Operator.TavilySearch]: TavilyForm,
|
||||
[Operator.TavilyExtract]: TavilyForm,
|
||||
[Operator.QueritSearch]: QueritForm,
|
||||
[Operator.WenCai]: WenCaiForm,
|
||||
[Operator.SearXNG]: SearXNGForm,
|
||||
[Operator.KeenableSearch]: KeenableForm,
|
||||
|
||||
41
web/src/pages/agent/form/tool-form/querit-form/index.tsx
Normal file
41
web/src/pages/agent/form/tool-form/querit-form/index.tsx
Normal file
@@ -0,0 +1,41 @@
|
||||
import { FormContainer } from '@/components/form-container';
|
||||
import { Form } from '@/components/ui/form';
|
||||
import { zodResolver } from '@hookform/resolvers/zod';
|
||||
import { memo, useMemo } from 'react';
|
||||
import { useForm } from 'react-hook-form';
|
||||
import { z } from 'zod';
|
||||
import { FormWrapper } from '../../components/form-wrapper';
|
||||
import { QueritFormWidgets } from '../../querit-form';
|
||||
import { useValues } from '../use-values';
|
||||
import { useQueritAgentFormChange } from './use-watch-change';
|
||||
import {
|
||||
QueritAgentFormSchema,
|
||||
deserializeQueritAgentFormValues,
|
||||
} from './utils';
|
||||
|
||||
function QueritForm() {
|
||||
const persistedValues = useValues();
|
||||
const values = useMemo(
|
||||
() => deserializeQueritAgentFormValues(persistedValues),
|
||||
[persistedValues],
|
||||
);
|
||||
const form = useForm<z.infer<typeof QueritAgentFormSchema>>({
|
||||
defaultValues: values,
|
||||
resolver: zodResolver(QueritAgentFormSchema),
|
||||
mode: 'onChange',
|
||||
});
|
||||
|
||||
useQueritAgentFormChange(form);
|
||||
|
||||
return (
|
||||
<Form {...form}>
|
||||
<FormWrapper>
|
||||
<FormContainer>
|
||||
<QueritFormWidgets includeQuery={false} />
|
||||
</FormContainer>
|
||||
</FormWrapper>
|
||||
</Form>
|
||||
);
|
||||
}
|
||||
|
||||
export default memo(QueritForm);
|
||||
@@ -0,0 +1,96 @@
|
||||
import { act, renderHook, waitFor } from '@testing-library/react';
|
||||
import { useForm } from 'react-hook-form';
|
||||
import { useQueritAgentFormChange } from './use-watch-change';
|
||||
|
||||
const mockUpdateAgentToolById = jest.fn();
|
||||
const agentNode = { id: 'agent:0' };
|
||||
|
||||
jest.mock('../../../store', () => ({
|
||||
__esModule: true,
|
||||
default: (selector: (state: Record<string, unknown>) => unknown) =>
|
||||
selector({
|
||||
clickedToolId: 'querit:0',
|
||||
clickedNodeId: 'agent:0',
|
||||
findUpstreamNodeById: () => agentNode,
|
||||
updateAgentToolById: mockUpdateAgentToolById,
|
||||
}),
|
||||
}));
|
||||
|
||||
const validValues = {
|
||||
api_key: '',
|
||||
count: 10,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
};
|
||||
|
||||
function useQueritAgentFormWatcher() {
|
||||
const form = useForm<Record<string, any>>({ defaultValues: validValues });
|
||||
useQueritAgentFormChange(form);
|
||||
return form;
|
||||
}
|
||||
|
||||
describe('useQueritAgentFormChange', () => {
|
||||
beforeEach(() => {
|
||||
mockUpdateAgentToolById.mockClear();
|
||||
});
|
||||
|
||||
it('does not write untouched initial values to the Agent node', async () => {
|
||||
renderHook(() => useQueritAgentFormWatcher());
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(mockUpdateAgentToolById).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('writes normalized valid edits without a query', async () => {
|
||||
const { result } = renderHook(() => useQueritAgentFormWatcher());
|
||||
|
||||
act(() => {
|
||||
result.current.setValue(
|
||||
'site_include',
|
||||
[{ value: ' docs.example.com ' }],
|
||||
{ shouldDirty: true },
|
||||
);
|
||||
result.current.setValue('count', '5', { shouldDirty: true });
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUpdateAgentToolById).toHaveBeenLastCalledWith(
|
||||
agentNode,
|
||||
'querit:0',
|
||||
{
|
||||
params: expect.objectContaining({
|
||||
count: 5,
|
||||
site_include: ['docs.example.com'],
|
||||
}),
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
const persistedParams =
|
||||
mockUpdateAgentToolById.mock.calls.at(-1)?.[2]?.params;
|
||||
expect(persistedParams).not.toHaveProperty('query');
|
||||
});
|
||||
|
||||
it('does not overwrite Agent tool params after an invalid edit', async () => {
|
||||
const { result } = renderHook(() => useQueritAgentFormWatcher());
|
||||
|
||||
act(() => {
|
||||
result.current.setValue('time_range', 'last week', {
|
||||
shouldDirty: true,
|
||||
});
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(mockUpdateAgentToolById).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,40 @@
|
||||
import { useEffect } from 'react';
|
||||
import { UseFormReturn, useWatch } from 'react-hook-form';
|
||||
import useGraphStore from '../../../store';
|
||||
import { validateQueritAgentFormValuesForPersistence } from './utils';
|
||||
|
||||
export function useQueritAgentFormChange(form?: UseFormReturn<any>) {
|
||||
const values = useWatch({ control: form?.control });
|
||||
const clickedToolId = useGraphStore((state) => state.clickedToolId);
|
||||
const clickedNodeId = useGraphStore((state) => state.clickedNodeId);
|
||||
const findUpstreamNodeById = useGraphStore(
|
||||
(state) => state.findUpstreamNodeById,
|
||||
);
|
||||
const updateAgentToolById = useGraphStore(
|
||||
(state) => state.updateAgentToolById,
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
const agentNode = findUpstreamNodeById(clickedNodeId);
|
||||
if (!agentNode || !form?.formState.isDirty) {
|
||||
return;
|
||||
}
|
||||
|
||||
const nextValues = validateQueritAgentFormValuesForPersistence(
|
||||
form.getValues(),
|
||||
);
|
||||
if (nextValues) {
|
||||
updateAgentToolById(agentNode, clickedToolId, {
|
||||
params: nextValues,
|
||||
});
|
||||
}
|
||||
}, [
|
||||
clickedNodeId,
|
||||
clickedToolId,
|
||||
findUpstreamNodeById,
|
||||
form,
|
||||
form?.formState.isDirty,
|
||||
updateAgentToolById,
|
||||
values,
|
||||
]);
|
||||
}
|
||||
66
web/src/pages/agent/form/tool-form/querit-form/utils.test.ts
Normal file
66
web/src/pages/agent/form/tool-form/querit-form/utils.test.ts
Normal file
@@ -0,0 +1,66 @@
|
||||
import {
|
||||
deserializeQueritAgentFormValues,
|
||||
validateQueritAgentFormValuesForPersistence,
|
||||
} from './utils';
|
||||
|
||||
describe('Querit Agent form persistence', () => {
|
||||
it('rehydrates persisted string arrays for the embedded form', () => {
|
||||
expect(
|
||||
deserializeQueritAgentFormValues({
|
||||
api_key: '',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
site_exclude: ['spam.example.com'],
|
||||
time_range: 'w1',
|
||||
country_include: ['united states'],
|
||||
language_include: ['en'],
|
||||
}),
|
||||
).toMatchObject({
|
||||
site_include: [{ value: 'docs.example.com' }],
|
||||
site_exclude: [{ value: 'spam.example.com' }],
|
||||
country_include: [{ value: 'united states' }],
|
||||
language_include: [{ value: 'en' }],
|
||||
});
|
||||
});
|
||||
|
||||
it('normalizes valid values without persisting a runtime query', () => {
|
||||
expect(
|
||||
validateQueritAgentFormValuesForPersistence({
|
||||
api_key: 'secret',
|
||||
query: 'must-not-be-persisted',
|
||||
count: '5',
|
||||
chunks_per_doc: '2',
|
||||
site_include: [{ value: ' docs.example.com ' }],
|
||||
site_exclude: [],
|
||||
time_range: 'w1',
|
||||
country_include: [{ value: ' united states ' }],
|
||||
language_include: [{ value: ' en ' }],
|
||||
}),
|
||||
).toEqual({
|
||||
api_key: 'secret',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
site_exclude: [],
|
||||
time_range: 'w1',
|
||||
country_include: ['united states'],
|
||||
language_include: ['en'],
|
||||
});
|
||||
});
|
||||
|
||||
it('does not produce persisted values when the embedded form is invalid', () => {
|
||||
expect(
|
||||
validateQueritAgentFormValuesForPersistence({
|
||||
api_key: '',
|
||||
count: 0,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
}),
|
||||
).toBeUndefined();
|
||||
});
|
||||
});
|
||||
22
web/src/pages/agent/form/tool-form/querit-form/utils.ts
Normal file
22
web/src/pages/agent/form/tool-form/querit-form/utils.ts
Normal file
@@ -0,0 +1,22 @@
|
||||
import {
|
||||
QueritFormSchema,
|
||||
deserializeQueritFormValues,
|
||||
serializeQueritFormValues,
|
||||
} from '../../querit-form/utils';
|
||||
|
||||
export const QueritAgentFormSchema = QueritFormSchema.omit({ query: true });
|
||||
|
||||
export function deserializeQueritAgentFormValues(values: Record<string, any>) {
|
||||
return deserializeQueritFormValues(values);
|
||||
}
|
||||
|
||||
export function validateQueritAgentFormValuesForPersistence(
|
||||
values: Record<string, any>,
|
||||
) {
|
||||
const result = QueritAgentFormSchema.safeParse(values);
|
||||
if (!result.success) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return serializeQueritFormValues(result.data);
|
||||
}
|
||||
53
web/src/pages/agent/hooks/clear-sensitive-fields.ts
Normal file
53
web/src/pages/agent/hooks/clear-sensitive-fields.ts
Normal file
@@ -0,0 +1,53 @@
|
||||
import { Operator } from '@/constants/agent';
|
||||
import { cloneDeepWith, get, isPlainObject } from 'lodash';
|
||||
|
||||
const apiKeyOperators = [
|
||||
Operator.TavilySearch,
|
||||
Operator.TavilyExtract,
|
||||
Operator.Google,
|
||||
Operator.KeenableSearch,
|
||||
Operator.BGPT,
|
||||
Operator.QueritSearch,
|
||||
];
|
||||
|
||||
function isQueritOperator(value: unknown) {
|
||||
if (typeof value !== 'string') {
|
||||
return false;
|
||||
}
|
||||
|
||||
return ['querit', 'queritsearch'].includes(
|
||||
value.replace(/_/g, '').toLowerCase(),
|
||||
);
|
||||
}
|
||||
|
||||
export function clearSensitiveFields<T>(obj: T): T {
|
||||
return cloneDeepWith(obj, (value) => {
|
||||
if (!isPlainObject(value)) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (
|
||||
(apiKeyOperators.includes(value.component_name) ||
|
||||
isQueritOperator(value.component_name)) &&
|
||||
get(value, 'params.api_key')
|
||||
) {
|
||||
return { ...value, params: { ...value.params, api_key: '' } };
|
||||
}
|
||||
|
||||
if (
|
||||
isQueritOperator(get(value, 'data.label')) &&
|
||||
get(value, 'data.form.api_key')
|
||||
) {
|
||||
return {
|
||||
...value,
|
||||
data: {
|
||||
...value.data,
|
||||
form: {
|
||||
...value.data.form,
|
||||
api_key: '',
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
});
|
||||
}
|
||||
7
web/src/pages/agent/hooks/querit-agent-initial-values.ts
Normal file
7
web/src/pages/agent/hooks/querit-agent-initial-values.ts
Normal file
@@ -0,0 +1,7 @@
|
||||
import { omit } from 'lodash';
|
||||
|
||||
export function getQueritAgentInitialValues(
|
||||
initialValues: Record<string, any>,
|
||||
) {
|
||||
return omit(initialValues, 'query', 'outputs');
|
||||
}
|
||||
@@ -38,6 +38,7 @@ import {
|
||||
initialParserValues,
|
||||
initialPubMedValues,
|
||||
initialBGPTValues,
|
||||
initialQueritValues,
|
||||
initialRetrievalValues,
|
||||
initialRewriteQuestionValues,
|
||||
initialSearXNGValues,
|
||||
@@ -167,6 +168,7 @@ export const useInitializeOperatorParams = () => {
|
||||
[Operator.Agent]: { ...initialAgentValues, llm_id: llmId },
|
||||
[Operator.Tool]: {},
|
||||
[Operator.TavilySearch]: initialTavilyValues,
|
||||
[Operator.QueritSearch]: initialQueritValues,
|
||||
[Operator.KeenableSearch]: initialKeenableValues,
|
||||
[Operator.UserFillUp]: initialUserFillUpValues,
|
||||
[Operator.StringTransform]: initialStringTransformValues,
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
import { getQueritAgentInitialValues } from './querit-agent-initial-values';
|
||||
|
||||
describe('getQueritAgentInitialValues', () => {
|
||||
it('creates Querit Agent params without runtime query or outputs', () => {
|
||||
expect(
|
||||
getQueritAgentInitialValues({
|
||||
api_key: '',
|
||||
query: 'sys.query',
|
||||
count: 10,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
outputs: {
|
||||
formalized_content: { value: '', type: 'string' },
|
||||
json: { value: {}, type: 'object' },
|
||||
},
|
||||
}),
|
||||
).toEqual({
|
||||
api_key: '',
|
||||
count: 10,
|
||||
chunks_per_doc: 3,
|
||||
site_include: [],
|
||||
site_exclude: [],
|
||||
time_range: '',
|
||||
country_include: [],
|
||||
language_include: [],
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,6 +1,7 @@
|
||||
import { omit, pick } from 'lodash';
|
||||
import { useCallback } from 'react';
|
||||
import { Operator } from '../constant';
|
||||
import { getQueritAgentInitialValues } from './querit-agent-initial-values';
|
||||
import { useInitializeOperatorParams } from './use-add-node';
|
||||
|
||||
export function useAgentToolInitialValues() {
|
||||
@@ -20,6 +21,8 @@ export function useAgentToolInitialValues() {
|
||||
return {
|
||||
api_key: '',
|
||||
};
|
||||
case Operator.QueritSearch:
|
||||
return getQueritAgentInitialValues(initialValues);
|
||||
case Operator.ExeSQL:
|
||||
return omit(initialValues, 'sql');
|
||||
case Operator.Bing:
|
||||
|
||||
126
web/src/pages/agent/hooks/use-export-json.test.ts
Normal file
126
web/src/pages/agent/hooks/use-export-json.test.ts
Normal file
@@ -0,0 +1,126 @@
|
||||
jest.mock('@/constants/agent', () => ({
|
||||
Operator: {
|
||||
TavilySearch: 'TavilySearch',
|
||||
TavilyExtract: 'TavilyExtract',
|
||||
Google: 'Google',
|
||||
KeenableSearch: 'KeenableSearch',
|
||||
BGPT: 'Bing',
|
||||
QueritSearch: 'QueritSearch',
|
||||
},
|
||||
}));
|
||||
|
||||
import { Operator } from '@/constants/agent';
|
||||
import { clearSensitiveFields } from './clear-sensitive-fields';
|
||||
|
||||
describe('clearSensitiveFields', () => {
|
||||
it('clears a Querit API key without removing non-sensitive params', () => {
|
||||
const dsl = {
|
||||
tools: [
|
||||
{
|
||||
component_name: Operator.QueritSearch,
|
||||
params: {
|
||||
api_key: 'querit-secret',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
expect(clearSensitiveFields(dsl)).toEqual({
|
||||
tools: [
|
||||
{
|
||||
component_name: Operator.QueritSearch,
|
||||
params: {
|
||||
api_key: '',
|
||||
count: 5,
|
||||
chunks_per_doc: 2,
|
||||
site_include: ['docs.example.com'],
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
expect(dsl.tools[0].params.api_key).toBe('querit-secret');
|
||||
});
|
||||
|
||||
it.each(['querit', 'querit_search', 'queritsearch'])(
|
||||
'clears a Querit API key for the %s registry alias',
|
||||
(componentName) => {
|
||||
const dsl = {
|
||||
tools: [
|
||||
{
|
||||
component_name: componentName,
|
||||
params: {
|
||||
api_key: 'querit-secret',
|
||||
count: 5,
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const sanitized = clearSensitiveFields(dsl);
|
||||
|
||||
expect(sanitized.tools[0].params.api_key).toBe('');
|
||||
expect(sanitized.tools[0].params.count).toBe(5);
|
||||
expect(dsl.tools[0].params.api_key).toBe('querit-secret');
|
||||
},
|
||||
);
|
||||
|
||||
it('clears a standalone Querit Canvas key from graph and components', () => {
|
||||
const dsl = {
|
||||
graph: {
|
||||
nodes: [
|
||||
{
|
||||
data: {
|
||||
label: Operator.QueritSearch,
|
||||
form: {
|
||||
api_key: 'graph-secret',
|
||||
count: 5,
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
components: {
|
||||
querit: {
|
||||
obj: {
|
||||
component_name: Operator.QueritSearch,
|
||||
params: {
|
||||
api_key: 'component-secret',
|
||||
count: 5,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const sanitized = clearSensitiveFields(dsl);
|
||||
|
||||
expect(sanitized.graph.nodes[0].data.form.api_key).toBe('');
|
||||
expect(sanitized.graph.nodes[0].data.form.count).toBe(5);
|
||||
expect(sanitized.components.querit.obj.params.api_key).toBe('');
|
||||
expect(sanitized.components.querit.obj.params.count).toBe(5);
|
||||
expect(dsl.graph.nodes[0].data.form.api_key).toBe('graph-secret');
|
||||
expect(dsl.components.querit.obj.params.api_key).toBe('component-secret');
|
||||
});
|
||||
|
||||
it('does not change standalone graph export behavior for other tools', () => {
|
||||
const dsl = {
|
||||
graph: {
|
||||
nodes: [
|
||||
{
|
||||
data: {
|
||||
label: Operator.TavilySearch,
|
||||
form: {
|
||||
api_key: 'existing-tavily-key',
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
|
||||
expect(clearSensitiveFields(dsl)).toEqual(dsl);
|
||||
});
|
||||
});
|
||||
@@ -1,30 +1,9 @@
|
||||
import { Operator } from '@/constants/agent';
|
||||
import { useFetchAgent } from '@/hooks/use-agent-request';
|
||||
import { downloadJsonFile } from '@/utils/file-util';
|
||||
import { cloneDeepWith, get, isPlainObject } from 'lodash';
|
||||
import { useCallback } from 'react';
|
||||
import useGraphStore from '../store';
|
||||
import { exportDsl } from '../utils/dsl-bridge';
|
||||
|
||||
/**
|
||||
* Recursively clear sensitive fields (api_key) from the DSL object
|
||||
*/
|
||||
const clearSensitiveFields = <T>(obj: T): T =>
|
||||
cloneDeepWith(obj, (value) => {
|
||||
if (
|
||||
isPlainObject(value) &&
|
||||
[
|
||||
Operator.TavilySearch,
|
||||
Operator.TavilyExtract,
|
||||
Operator.Google,
|
||||
Operator.KeenableSearch,
|
||||
Operator.BGPT,
|
||||
].includes(value.component_name) &&
|
||||
get(value, 'params.api_key')
|
||||
) {
|
||||
return { ...value, params: { ...value.params, api_key: '' } };
|
||||
}
|
||||
});
|
||||
import { clearSensitiveFields } from './clear-sensitive-fields';
|
||||
|
||||
export const useHandleExportJsonFile = () => {
|
||||
const { data } = useFetchAgent();
|
||||
|
||||
24
web/src/pages/agent/log-sheet/tool-name.test.ts
Normal file
24
web/src/pages/agent/log-sheet/tool-name.test.ts
Normal file
@@ -0,0 +1,24 @@
|
||||
jest.mock('@/constants/agent', () => ({
|
||||
Operator: {
|
||||
QueritSearch: 'QueritSearch',
|
||||
},
|
||||
}));
|
||||
|
||||
import { Operator } from '@/constants/agent';
|
||||
import { getToolOperatorName } from './tool-name';
|
||||
|
||||
describe('getToolOperatorName', () => {
|
||||
it.each(['QueritSearch', 'querit_search'])(
|
||||
'maps the Querit timeline name %p to its operator',
|
||||
(toolName) => {
|
||||
expect(getToolOperatorName(toolName)).toBe(Operator.QueritSearch);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([undefined, null, ''])(
|
||||
'returns an empty name for the missing value %p',
|
||||
(toolName) => {
|
||||
expect(getToolOperatorName(toolName)).toBe('');
|
||||
},
|
||||
);
|
||||
});
|
||||
17
web/src/pages/agent/log-sheet/tool-name.ts
Normal file
17
web/src/pages/agent/log-sheet/tool-name.ts
Normal file
@@ -0,0 +1,17 @@
|
||||
import { Operator } from '@/constants/agent';
|
||||
|
||||
export function getToolOperatorName(toolName?: string | null) {
|
||||
if (!toolName) {
|
||||
return '';
|
||||
}
|
||||
|
||||
const normalizedName = toolName.replaceAll('_', '').toLowerCase();
|
||||
if (normalizedName === Operator.QueritSearch.toLowerCase()) {
|
||||
return Operator.QueritSearch;
|
||||
}
|
||||
|
||||
return toolName
|
||||
.split('_')
|
||||
.map((word) => word.charAt(0).toUpperCase() + word.slice(1).toLowerCase())
|
||||
.join('');
|
||||
}
|
||||
@@ -16,6 +16,7 @@ import { isEmpty } from 'lodash';
|
||||
import { Operator } from '../constant';
|
||||
import { JsonViewer } from '../form/components/json-viewer';
|
||||
import OperatorIcon, { SVGIconMap } from '../operator-icon';
|
||||
import { getToolOperatorName } from './tool-name';
|
||||
import { toLowerCaseStringAndDeleteChar, typeMap } from './workflow-timeline';
|
||||
type IToolIcon =
|
||||
| Operator.ArXiv
|
||||
@@ -28,6 +29,7 @@ type IToolIcon =
|
||||
| Operator.BGPT
|
||||
| Operator.TavilyExtract
|
||||
| Operator.TavilySearch
|
||||
| Operator.QueritSearch
|
||||
| Operator.KeenableSearch
|
||||
| Operator.Wikipedia
|
||||
| Operator.YahooFinance
|
||||
@@ -73,7 +75,7 @@ const ToolTimelineItem = ({
|
||||
return (
|
||||
<>
|
||||
{filteredTools?.map((tool, idx) => {
|
||||
const toolName = capitalizeWords(tool.tool_name, '_').join('');
|
||||
const toolName = getToolOperatorName(tool.tool_name);
|
||||
|
||||
return (
|
||||
<TimelineItem
|
||||
|
||||
@@ -51,6 +51,7 @@ export const typeMap = {
|
||||
textProcessing: t('flow.logTimeline.textProcessing'),
|
||||
tavilySearch: t('flow.logTimeline.tavilySearch'),
|
||||
tavilyExtract: t('flow.logTimeline.tavilyExtract'),
|
||||
queritSearch: t('flow.logTimeline.queritSearch'),
|
||||
exeSQL: t('flow.logTimeline.exeSQL'),
|
||||
google: t('flow.logTimeline.google'),
|
||||
duckDuckGo: t('flow.logTimeline.google'),
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { IconFontFill } from '@/components/icon-font';
|
||||
import SvgIcon from '@/components/svg-icon';
|
||||
import queritLogo from '@/assets/querit.png';
|
||||
import { cn } from '@/lib/utils';
|
||||
import {
|
||||
Columns3Cog,
|
||||
@@ -51,6 +52,7 @@ export const SVGIconMap = {
|
||||
[Operator.KeenableSearch]: 'keenable',
|
||||
[Operator.TavilyExtract]: 'tavily',
|
||||
[Operator.TavilySearch]: 'tavily',
|
||||
[Operator.QueritSearch]: 'querit',
|
||||
[Operator.Wikipedia]: 'wikipedia',
|
||||
[Operator.YahooFinance]: 'yahoo-finance',
|
||||
[Operator.WenCai]: 'wencai',
|
||||
@@ -93,6 +95,16 @@ const OperatorIcon = ({ name, className }: IProps) => {
|
||||
const svgIcon = SVGIconMap[name as keyof typeof SVGIconMap];
|
||||
const LucideIcon = LucideIconMap[name as keyof typeof LucideIconMap];
|
||||
|
||||
if (name === Operator.QueritSearch) {
|
||||
return (
|
||||
<img
|
||||
src={queritLogo}
|
||||
alt=""
|
||||
className={cn('size-5 object-contain', className)}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
if (name === Operator.Begin) {
|
||||
return (
|
||||
<div
|
||||
|
||||
@@ -51,6 +51,7 @@ export const typeMap = {
|
||||
textProcessing: t('flow.logTimeline.textProcessing'),
|
||||
tavilySearch: t('flow.logTimeline.tavilySearch'),
|
||||
tavilyExtract: t('flow.logTimeline.tavilyExtract'),
|
||||
queritSearch: t('flow.logTimeline.queritSearch'),
|
||||
exeSQL: t('flow.logTimeline.exeSQL'),
|
||||
google: t('flow.logTimeline.google'),
|
||||
duckDuckGo: t('flow.logTimeline.google'),
|
||||
|
||||
Reference in New Issue
Block a user