mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-02 16:55:42 +08:00
fix(agent): filter TuShare news with upstream keyword input (#16361)
## Summary TuShare required non-empty upstream input but filtered fetched news with the static `keyword` param (default empty string), so agent-provided keywords were ignored. Use `self._param.keyword or ans` when filtering, matching how AkShare uses upstream input for its query. Fixes #16360 ## Test plan - [x] `test_tushare_filters_with_upstream_keyword_when_param_empty` mocks the API and asserts only rows matching the upstream keyword are returned --------- Co-authored-by: yzc <yuzhichang@gmail.com> Co-authored-by: Harsh Kashyap <harshkashyap@Harshs-MacBook-Pro.local>
This commit is contained in:
@@ -14,6 +14,7 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
import json
|
||||
import logging
|
||||
from abc import ABC
|
||||
import pandas as pd
|
||||
import time
|
||||
@@ -73,7 +74,18 @@ class TuShare(ComponentBase, ABC):
|
||||
df.columns = response['data']['fields']
|
||||
if self.check_if_canceled("TuShare processing"):
|
||||
return
|
||||
tus_res.append({"content": (df[df['content'].str.contains(self._param.keyword, case=False)]).to_markdown()})
|
||||
keyword = self._param.keyword or ans
|
||||
logging.info(
|
||||
"TuShare news filter keyword source=%s",
|
||||
"param.keyword" if self._param.keyword else "upstream_input",
|
||||
)
|
||||
tus_res.append(
|
||||
{
|
||||
"content": (
|
||||
df[df["content"].str.contains(keyword, case=False, na=False, regex=False)]
|
||||
).to_markdown()
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
if self.check_if_canceled("TuShare processing"):
|
||||
return
|
||||
|
||||
91
test/unit_test/agent/tools/test_tushare.py
Normal file
91
test/unit_test/agent/tools/test_tushare.py
Normal file
@@ -0,0 +1,91 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agent.tools.tushare import TuShare, TuShareParam
|
||||
|
||||
|
||||
class _Canvas:
|
||||
def is_canceled(self):
|
||||
return False
|
||||
|
||||
|
||||
def _tushare(param=None):
|
||||
cpn = TuShare.__new__(TuShare)
|
||||
cpn._canvas = _Canvas()
|
||||
cpn._param = param or TuShareParam()
|
||||
return cpn
|
||||
|
||||
|
||||
def _response():
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"items": [
|
||||
[1, "Apple earnings beat expectations"],
|
||||
[2, "Google cloud revenue grows"],
|
||||
[3, "C++ regex special chars should not break filtering"],
|
||||
],
|
||||
"fields": ["id", "content"],
|
||||
},
|
||||
}
|
||||
return response
|
||||
|
||||
|
||||
@pytest.mark.p1
|
||||
def test_tushare_filters_with_upstream_keyword_when_param_empty():
|
||||
cpn = _tushare()
|
||||
|
||||
with patch.object(TuShare, "get_input", return_value={"content": ["Apple"]}):
|
||||
with patch("agent.tools.tushare.requests.post", return_value=_response()):
|
||||
result = cpn._run([])
|
||||
|
||||
text = result.iloc[0]["content"]
|
||||
assert "Apple" in text
|
||||
assert "Google" not in text
|
||||
|
||||
|
||||
@pytest.mark.p1
|
||||
def test_tushare_prefers_explicit_param_keyword_over_upstream_input():
|
||||
param = TuShareParam()
|
||||
param.keyword = "Google"
|
||||
cpn = _tushare(param)
|
||||
|
||||
with patch.object(TuShare, "get_input", return_value={"content": ["Apple"]}):
|
||||
with patch("agent.tools.tushare.requests.post", return_value=_response()):
|
||||
result = cpn._run([])
|
||||
|
||||
text = result.iloc[0]["content"]
|
||||
assert "Google" in text
|
||||
assert "Apple" not in text
|
||||
|
||||
|
||||
@pytest.mark.p1
|
||||
def test_tushare_treats_keyword_as_literal_text():
|
||||
param = TuShareParam()
|
||||
param.keyword = "C++"
|
||||
cpn = _tushare(param)
|
||||
|
||||
with patch.object(TuShare, "get_input", return_value={"content": ["ignored"]}):
|
||||
with patch("agent.tools.tushare.requests.post", return_value=_response()):
|
||||
result = cpn._run([])
|
||||
|
||||
text = result.iloc[0]["content"]
|
||||
assert "C++ regex special chars should not break filtering" in text
|
||||
Reference in New Issue
Block a user