Files
2026-07-13 13:26:09 +08:00

125 lines
5.4 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import os
import json
import requests
from typing import Union, List
from qwen_agent.tools.base import BaseTool, register_tool
from concurrent.futures import ThreadPoolExecutor
SELECT_URL_PROMPT = """Please process the following search results and user goal to extract all relevant urls:
## **Search Results**
{search_results}
## **User Goal**
{goal}
## **Task Guidelines**
1. **Content Scanning for Rational**: Locate all the relevant **urls** directly related to the user's goal according to the titles, snippets, and url format in the search content.
2. **Extract relevant urls for goals**: Identify and extract all the **relevant urls** from the content, you never miss any important urls, output the **full original urls** of the content as far as possible. Ensure the urls are valid and complete.
**Final Output Format using JSON format has "rational", "urls" feilds**
Output example:
```json{{
"rational": "The rational is...",
"urls": ["url1", "url2", "url3"]
}}
'''
"""
@register_tool('select_related_url')
class SelectURL(BaseTool):
def call_server_align(self, msgs, max_retries=2):
url_llm = "http://localhost:6002/v1/chat/completions"
headers = {
'Content-Type': 'application/json',
'Authorization': "EMPTY"
}
payload = json.dumps({
"model": os.environ.get("SUMMARY_MODEL_PATH"),
"temperature": 0.7,
"top_p": 0.8,
"max_tokens": 80000, ## 16000
"presence_penalty": 1.5,
"messages": msgs,
"chat_template_kwargs": {
"enable_thinking": False
}
})
response = None
for _ in range(max_retries):
try:
response = requests.request("POST", url_llm, headers=headers, data=payload, timeout=240)
data = response.json()
if "error" in data and data["error"]["message"] == "Provider returned error":
print("error in message")
return ""
content = data["choices"][0]["message"]["content"].replace("```json", "").replace("```", "")
try:
content = json.loads(content)
return content
except:
print("Error happens in url selection. Retry again.")
left_pos, right_pos = content.index("{"), content.rindex("}")
if left_pos > 0 and right_pos > left_pos:
content = content[left_pos:right_pos+1]
return json.loads(content)
return content
except Exception as e:
print("URL Select Error, retry again", e)
continue
# print(response.text)
return ""
def call(self, params: Union[str, dict], **kwargs) -> str:
try:
params = self._verify_json_format_args(params)
search_results = params["search_results"]
goal = params["goal"]
except:
return "[Search] Invalid request format: Input must be a JSON object containing 'query' field"
messages = [{"role":"user","content": SELECT_URL_PROMPT.format(search_results=search_results, goal=goal)}]
select_url_json = self.call_server_align(messages)
return select_url_json
if __name__ == "__main__":
# tool = Search()
search_content = f"""
## Web Results
1. [The Hong Kong University of Science and Technology World's ...](https://hkust.edu.hk/)
Hong Kong University of Science and Technology (HKUST) is a world-class international research university, leading research excellence in science, ...
2. [Hong Kong University of Science and Technology - Wikipedia](https://en.wikipedia.org/wiki/Hong_Kong_University_of_Science_and_Technology)
HKUST is a public research university in Sai Kung District, New Territories, Hong Kong. Founded in 1991, it was the territory's third institution to be granted ...
3. [The Hong Kong University of Science and Technology](https://www.topuniversities.com/universities/hong-kong-university-science-technology)
We offer interdisciplinary programs across science, technology, engineering, business, and humanities, complemented by a strong foundation in AI, data science, ...
4. [Hong Kong University of Science and Technology (HKUST)](https://experience.cornell.edu/opportunities/hong-kong-university-science-and-technology-hkust)
HKUST is a leading research university with English courses, guaranteed housing, and a diverse student body. It is a Cornell Global Hub partner with a 150-acre ...
5. [The Hong Kong University of Science and Technology (Guangzhou)](https://www.hkust-gz.edu.cn/)
HKUST(GZ) will explore new frontiers in cross-disciplinary education and innovate pedagogies. By doing so, HKUST(GZ) will play a constructive role in ...
6. [The Hong Kong University of Science and Technology (@hkust)](https://www.instagram.com/hkust/?hl=en)
The Hong Kong University of Science and Technology (HKUST) Official Account. Tag us with #HKUST to let us repost your content on our channels!
"""
# print(tool.call({"query": ["When Do LLMs Help With Node Classification? A Comprehensive Analysis", "CUHK DB Group"]}))
select_url = SelectURL()
# print(select_url.call({"search_results": search_content, "goal": "introduce the university HKUST"}))
print(select_url.call({"search_results": "hello", "goal": "introduce the university HKUST"}))