125 lines
5.4 KiB
Python
125 lines
5.4 KiB
Python
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"}))
|