diff --git a/agent/tools/querit.py b/agent/tools/querit.py new file mode 100644 index 0000000000..6725342a94 --- /dev/null +++ b/agent/tools/querit.py @@ -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) diff --git a/internal/agent/component/tool_component.go b/internal/agent/component/tool_component.go index f4b5919a36..2eaca204b5 100644 --- a/internal/agent/component/tool_component.go +++ b/internal/agent/component/tool_component.go @@ -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"}, diff --git a/internal/agent/component/tool_component_test.go b/internal/agent/component/tool_component_test.go index c2b0b9caf4..64decfa181 100644 --- a/internal/agent/component/tool_component_test.go +++ b/internal/agent/component/tool_component_test.go @@ -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) { diff --git a/internal/agent/tool/querit.go b/internal/agent/tool/querit.go new file mode 100644 index 0000000000..5555c97f6e --- /dev/null +++ b/internal/agent/tool/querit.go @@ -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) +} diff --git a/internal/agent/tool/querit_test.go b/internal/agent/tool/querit_test.go new file mode 100644 index 0000000000..97d87b1a8d --- /dev/null +++ b/internal/agent/tool/querit_test.go @@ -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": "![img](data:image/png;base64,AAAA)"}, + map[string]any{"title": "cleaned", "snippet": "before ![img](data:image/png;base64,AAAA) 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) +} diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index b64e209291..ca43ca7830 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -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 := "" diff --git a/internal/agent/tool/registry_test.go b/internal/agent/tool/registry_test.go index 03b5fb8118..27b695fbef 100644 --- a/internal/agent/tool/registry_test.go +++ b/internal/agent/tool/registry_test.go @@ -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] diff --git a/internal/agent/tool/tool2component.go b/internal/agent/tool/tool2component.go index df487921d7..f637fc0969 100644 --- a/internal/agent/tool/tool2component.go +++ b/internal/agent/tool/tool2component.go @@ -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 diff --git a/internal/common/environments.go b/internal/common/environments.go index 8f28502c38..1e0fd91db4 100644 --- a/internal/common/environments.go +++ b/internal/common/environments.go @@ -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" diff --git a/test/unit_test/agent/tools/test_querit_unit.py b/test/unit_test/agent/tools/test_querit_unit.py new file mode 100644 index 0000000000..3c88cdda43 --- /dev/null +++ b/test/unit_test/agent/tools/test_querit_unit.py @@ -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"] diff --git a/tools/gen-component-parity/main.go b/tools/gen-component-parity/main.go index fbe7ab392f..8c644f283d 100644 --- a/tools/gen-component-parity/main.go +++ b/tools/gen-component-parity/main.go @@ -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", } diff --git a/web/src/assets/querit.png b/web/src/assets/querit.png new file mode 100644 index 0000000000..223524516a Binary files /dev/null and b/web/src/assets/querit.png differ diff --git a/web/src/constants/agent.tsx b/web/src/constants/agent.tsx index 5a2b3f9565..6eb5dcce0e 100644 --- a/web/src/constants/agent.tsx +++ b/web/src/constants/agent.tsx @@ -95,6 +95,7 @@ export enum Operator { Tool = 'Tool', TavilySearch = 'TavilySearch', TavilyExtract = 'TavilyExtract', + QueritSearch = 'QueritSearch', UserFillUp = 'UserFillUp', StringTransform = 'StringTransform', SearXNG = 'SearXNG', diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts index 43647640ef..5d4b6d418f 100644 --- a/web/src/locales/en.ts +++ b/web/src/locales/en.ts @@ -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', diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts index d5c82ba222..97923e30f5 100644 --- a/web/src/locales/zh.ts +++ b/web/src/locales/zh.ts @@ -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: '正在网上搜索', diff --git a/web/src/pages/agent/canvas/node/dropdown/accordion-operators.tsx b/web/src/pages/agent/canvas/node/dropdown/accordion-operators.tsx index b19b4a7b6d..a4a6cdd9e5 100644 --- a/web/src/pages/agent/canvas/node/dropdown/accordion-operators.tsx +++ b/web/src/pages/agent/canvas/node/dropdown/accordion-operators.tsx @@ -109,6 +109,7 @@ export function AccordionOperators({ operators={[ Operator.TavilySearch, Operator.TavilyExtract, + Operator.QueritSearch, Operator.ExeSQL, Operator.Google, Operator.YahooFinance, diff --git a/web/src/pages/agent/constant/index.tsx b/web/src/pages/agent/constant/index.tsx index 5f1e4050bc..15242fa3d1 100644 --- a/web/src/pages/agent/constant/index.tsx +++ b/web/src/pages/agent/constant/index.tsx @@ -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', diff --git a/web/src/pages/agent/form-sheet/form-config-map.tsx b/web/src/pages/agent/form-sheet/form-config-map.tsx index 7e6b12c678..5d7261dbe1 100644 --- a/web/src/pages/agent/form-sheet/form-config-map.tsx +++ b/web/src/pages/agent/form-sheet/form-config-map.tsx @@ -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, }, diff --git a/web/src/pages/agent/form/agent-form/tool-popover/tool-command.tsx b/web/src/pages/agent/form/agent-form/tool-popover/tool-command.tsx index 400bba78f5..7ed2af7b87 100644 --- a/web/src/pages/agent/form/agent-form/tool-popover/tool-command.tsx +++ b/web/src/pages/agent/form/agent-form/tool-popover/tool-command.tsx @@ -23,6 +23,7 @@ const Menus = [ list: [ Operator.TavilySearch, Operator.TavilyExtract, + Operator.QueritSearch, Operator.Google, // Operator.Bing, Operator.DuckDuckGo, diff --git a/web/src/pages/agent/form/querit-form/dynamic-list.tsx b/web/src/pages/agent/form/querit-form/dynamic-list.tsx new file mode 100644 index 0000000000..197e8e332c --- /dev/null +++ b/web/src/pages/agent/form/querit-form/dynamic-list.tsx @@ -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 ( + + {label} +
+ {fields.map((item, index) => ( +
+ ( + + + + + + + )} + /> + +
+ ))} +
+ append({ value: '' })}> + {t('common.add')} + + + ); +} diff --git a/web/src/pages/agent/form/querit-form/index.tsx b/web/src/pages/agent/form/querit-form/index.tsx new file mode 100644 index 0000000000..1fb5e4aa50 --- /dev/null +++ b/web/src/pages/agent/form/querit-form/index.tsx @@ -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 ( + <> + + {includeQuery && } + ( + + + {t('queritCount')} + + + + + + + )} + /> + ( + + + {t('queritChunksPerDoc')} + + + + + + + )} + /> + + + ( + + + {t('queritTimeRange')} + + + + + + + )} + /> + + + + ); +} + +function QueritForm({ node }: INextOperatorForm) { + const values = useValues(node); + const form = useForm>({ + defaultValues: values, + resolver: zodResolver(QueritFormSchema), + mode: 'onChange', + }); + + useWatchFormChange(node?.id, form); + + return ( +
+ + + + + +
+ +
+
+ ); +} + +export default memo(QueritForm); diff --git a/web/src/pages/agent/form/querit-form/use-values.ts b/web/src/pages/agent/form/querit-form/use-values.ts new file mode 100644 index 0000000000..8d77013601 --- /dev/null +++ b/web/src/pages/agent/form/querit-form/use-values.ts @@ -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]); +} diff --git a/web/src/pages/agent/form/querit-form/use-watch-change.test.tsx b/web/src/pages/agent/form/querit-form/use-watch-change.test.tsx new file mode 100644 index 0000000000..d6e9cd62f2 --- /dev/null +++ b/web/src/pages/agent/form/querit-form/use-watch-change.test.tsx @@ -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) => 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) { + 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(); + }); +}); diff --git a/web/src/pages/agent/form/querit-form/use-watch-change.ts b/web/src/pages/agent/form/querit-form/use-watch-change.ts new file mode 100644 index 0000000000..e494a9289a --- /dev/null +++ b/web/src/pages/agent/form/querit-form/use-watch-change.ts @@ -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) { + 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]); +} diff --git a/web/src/pages/agent/form/querit-form/utils.test.ts b/web/src/pages/agent/form/querit-form/utils.test.ts new file mode 100644 index 0000000000..615cec4c93 --- /dev/null +++ b/web/src/pages/agent/form/querit-form/utils.test.ts @@ -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: [], + }); + }); +}); diff --git a/web/src/pages/agent/form/querit-form/utils.ts b/web/src/pages/agent/form/querit-form/utils.ts new file mode 100644 index 0000000000..a25cf4a1aa --- /dev/null +++ b/web/src/pages/agent/form/querit-form/utils.ts @@ -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; + +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, + }); +} diff --git a/web/src/pages/agent/form/tool-form/constant.tsx b/web/src/pages/agent/form/tool-form/constant.tsx index 34621adb22..772ed10c6a 100644 --- a/web/src/pages/agent/form/tool-form/constant.tsx +++ b/web/src/pages/agent/form/tool-form/constant.tsx @@ -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, diff --git a/web/src/pages/agent/form/tool-form/querit-form/index.tsx b/web/src/pages/agent/form/tool-form/querit-form/index.tsx new file mode 100644 index 0000000000..42df1ed88f --- /dev/null +++ b/web/src/pages/agent/form/tool-form/querit-form/index.tsx @@ -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>({ + defaultValues: values, + resolver: zodResolver(QueritAgentFormSchema), + mode: 'onChange', + }); + + useQueritAgentFormChange(form); + + return ( +
+ + + + + +
+ ); +} + +export default memo(QueritForm); diff --git a/web/src/pages/agent/form/tool-form/querit-form/use-watch-change.test.tsx b/web/src/pages/agent/form/tool-form/querit-form/use-watch-change.test.tsx new file mode 100644 index 0000000000..4a1b83d122 --- /dev/null +++ b/web/src/pages/agent/form/tool-form/querit-form/use-watch-change.test.tsx @@ -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) => 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>({ 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(); + }); +}); diff --git a/web/src/pages/agent/form/tool-form/querit-form/use-watch-change.ts b/web/src/pages/agent/form/tool-form/querit-form/use-watch-change.ts new file mode 100644 index 0000000000..e24213f5d7 --- /dev/null +++ b/web/src/pages/agent/form/tool-form/querit-form/use-watch-change.ts @@ -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) { + 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, + ]); +} diff --git a/web/src/pages/agent/form/tool-form/querit-form/utils.test.ts b/web/src/pages/agent/form/tool-form/querit-form/utils.test.ts new file mode 100644 index 0000000000..9c0e69669b --- /dev/null +++ b/web/src/pages/agent/form/tool-form/querit-form/utils.test.ts @@ -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(); + }); +}); diff --git a/web/src/pages/agent/form/tool-form/querit-form/utils.ts b/web/src/pages/agent/form/tool-form/querit-form/utils.ts new file mode 100644 index 0000000000..29c106cf87 --- /dev/null +++ b/web/src/pages/agent/form/tool-form/querit-form/utils.ts @@ -0,0 +1,22 @@ +import { + QueritFormSchema, + deserializeQueritFormValues, + serializeQueritFormValues, +} from '../../querit-form/utils'; + +export const QueritAgentFormSchema = QueritFormSchema.omit({ query: true }); + +export function deserializeQueritAgentFormValues(values: Record) { + return deserializeQueritFormValues(values); +} + +export function validateQueritAgentFormValuesForPersistence( + values: Record, +) { + const result = QueritAgentFormSchema.safeParse(values); + if (!result.success) { + return undefined; + } + + return serializeQueritFormValues(result.data); +} diff --git a/web/src/pages/agent/hooks/clear-sensitive-fields.ts b/web/src/pages/agent/hooks/clear-sensitive-fields.ts new file mode 100644 index 0000000000..2129cd39f0 --- /dev/null +++ b/web/src/pages/agent/hooks/clear-sensitive-fields.ts @@ -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(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: '', + }, + }, + }; + } + }); +} diff --git a/web/src/pages/agent/hooks/querit-agent-initial-values.ts b/web/src/pages/agent/hooks/querit-agent-initial-values.ts new file mode 100644 index 0000000000..b54a62da59 --- /dev/null +++ b/web/src/pages/agent/hooks/querit-agent-initial-values.ts @@ -0,0 +1,7 @@ +import { omit } from 'lodash'; + +export function getQueritAgentInitialValues( + initialValues: Record, +) { + return omit(initialValues, 'query', 'outputs'); +} diff --git a/web/src/pages/agent/hooks/use-add-node.ts b/web/src/pages/agent/hooks/use-add-node.ts index 4d7e6fe9a9..6bb13d9ae7 100644 --- a/web/src/pages/agent/hooks/use-add-node.ts +++ b/web/src/pages/agent/hooks/use-add-node.ts @@ -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, diff --git a/web/src/pages/agent/hooks/use-agent-tool-initial-values.test.ts b/web/src/pages/agent/hooks/use-agent-tool-initial-values.test.ts new file mode 100644 index 0000000000..02a0422d6b --- /dev/null +++ b/web/src/pages/agent/hooks/use-agent-tool-initial-values.test.ts @@ -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: [], + }); + }); +}); diff --git a/web/src/pages/agent/hooks/use-agent-tool-initial-values.ts b/web/src/pages/agent/hooks/use-agent-tool-initial-values.ts index 3d4b8112e0..3250a55fde 100644 --- a/web/src/pages/agent/hooks/use-agent-tool-initial-values.ts +++ b/web/src/pages/agent/hooks/use-agent-tool-initial-values.ts @@ -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: diff --git a/web/src/pages/agent/hooks/use-export-json.test.ts b/web/src/pages/agent/hooks/use-export-json.test.ts new file mode 100644 index 0000000000..1fd1a4fd1e --- /dev/null +++ b/web/src/pages/agent/hooks/use-export-json.test.ts @@ -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); + }); +}); diff --git a/web/src/pages/agent/hooks/use-export-json.ts b/web/src/pages/agent/hooks/use-export-json.ts index ff9191aa84..9f2e312159 100644 --- a/web/src/pages/agent/hooks/use-export-json.ts +++ b/web/src/pages/agent/hooks/use-export-json.ts @@ -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 = (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(); diff --git a/web/src/pages/agent/log-sheet/tool-name.test.ts b/web/src/pages/agent/log-sheet/tool-name.test.ts new file mode 100644 index 0000000000..2f99622724 --- /dev/null +++ b/web/src/pages/agent/log-sheet/tool-name.test.ts @@ -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(''); + }, + ); +}); diff --git a/web/src/pages/agent/log-sheet/tool-name.ts b/web/src/pages/agent/log-sheet/tool-name.ts new file mode 100644 index 0000000000..80e87f2c15 --- /dev/null +++ b/web/src/pages/agent/log-sheet/tool-name.ts @@ -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(''); +} diff --git a/web/src/pages/agent/log-sheet/tool-timeline-item.tsx b/web/src/pages/agent/log-sheet/tool-timeline-item.tsx index c40741e14e..8ea32586b2 100644 --- a/web/src/pages/agent/log-sheet/tool-timeline-item.tsx +++ b/web/src/pages/agent/log-sheet/tool-timeline-item.tsx @@ -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 ( { const svgIcon = SVGIconMap[name as keyof typeof SVGIconMap]; const LucideIcon = LucideIconMap[name as keyof typeof LucideIconMap]; + if (name === Operator.QueritSearch) { + return ( + + ); + } + if (name === Operator.Begin) { return (