wehub-resource-sync
243 lines
9.1 KiB
Python
243 lines
9.1 KiB
Python
import json
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import List, Union
|
|
import requests
|
|
from qwen_agent.tools.base import BaseTool, register_tool
|
|
# from prompt import EXTRACTOR_PROMPT
|
|
import os
|
|
from openai import OpenAI
|
|
import random
|
|
from topsdk.client import TopApiClient,TopException
|
|
from topsdk.defaultability.defaultability import Defaultability
|
|
from topsdk.defaultability.request.alibaba_aidata_aignite_application_run_request import AlibabaAidataAigniteApplicationRunAigniteApplicationExecuteReqDTO,AlibabaAidataAigniteApplicationRunRequest
|
|
from urllib.parse import urlparse, unquote
|
|
import time
|
|
|
|
import litellm
|
|
from qwen_agent.llm.schema import Message
|
|
from qwen_agent.utils.utils import build_text_completion_prompt
|
|
import tiktoken
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
def count_tokens(messages, model="gpt-4o"):
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained(self.llm_local_path)
|
|
except Exception as e:
|
|
tokenizer = tiktoken.encoding_for_model(model)
|
|
|
|
if isinstance(messages, list):
|
|
full_message = [Message(**x) for x in messages]
|
|
elif isinstance(messages, str):
|
|
return len(tokenizer.encode(messages))
|
|
else:
|
|
raise ValueError("Invalid message type")
|
|
full_prompt = build_text_completion_prompt(full_message, allow_special=True)
|
|
|
|
return len(tokenizer.encode(full_prompt))
|
|
|
|
|
|
class ScraperReader:
|
|
"""Scraper内容读取器"""
|
|
|
|
def __init__(self):
|
|
self.scraper_key = os.getenv("SCRAPER_API_KEY")
|
|
self.scraper_country_code = "us"
|
|
self.scraper_time_out = 40
|
|
def readpage(self, url: str) -> str:
|
|
payload = {'api_key': self.scraper_key,
|
|
'url': url,
|
|
'output_format': 'markdown',
|
|
'country_code': self.scraper_country_code }
|
|
max_retries = 2
|
|
for attempt in range(max_retries):
|
|
try:
|
|
r = requests.get('https://api.scraperapi.com/', params=payload, timeout=self.scraper_time_out)
|
|
content = r.text
|
|
return content
|
|
except requests.exceptions.Timeout:
|
|
# 超时情况下返回默认内容
|
|
content = "[visit] Request timed out after {} seconds. Using default content.".format(self.scraper_time_out)
|
|
except Exception as e:
|
|
content ="[visit] Failed to read page."
|
|
return content
|
|
|
|
|
|
@register_tool('retrieve', allow_overwrite=True)
|
|
class Retrieve(BaseTool):
|
|
"""
|
|
Retrieve webpages associated with the given id and return a summary.
|
|
The mapping between URL and id is supplied via the url2id dictionary.
|
|
"""
|
|
name = "retrieve"
|
|
description = "Read the webpage(s) whose id matches the given id and return the summary."
|
|
|
|
parameters = {
|
|
"type": "object",
|
|
"properties": {
|
|
"url_id": {
|
|
"type": ["string", "array"],
|
|
"items": {
|
|
"type": "string"
|
|
},
|
|
"minItems": 1,
|
|
"description": "The URL ID(s) of the webpage(s) to visit. Can be a single URL ID or an array of URL IDs."
|
|
},
|
|
"goal": {
|
|
"type": "string",
|
|
"description": "The goal of the visit for webpage(s)."
|
|
}
|
|
},
|
|
"required": ["url_id", "goal"]
|
|
}
|
|
|
|
# ---------------------------
|
|
# Main entry point
|
|
# ---------------------------
|
|
def call(self, params: Union[str, dict], **kwargs) -> str:
|
|
"""
|
|
params must contain:
|
|
• url2id : dict mapping URL → id
|
|
• url_id : target id
|
|
• goal : reading goal (passed through to readpage)
|
|
if "url2page" is present, it is used to retrieve the page content
|
|
"""
|
|
max_tokens = 20000
|
|
try:
|
|
url2id = params["url2id"]
|
|
target_id = params["url_id"]
|
|
if isinstance(target_id, list):
|
|
target_id = [int(id_.replace("id_", "")) for id_ in target_id]
|
|
elif isinstance(target_id, str):
|
|
target_id = [int(target_id.replace("id_", ""))]
|
|
else:
|
|
raise ValueError("Invalid url_id type")
|
|
goal = params["goal"]
|
|
assert isinstance(url2id, dict)
|
|
except Exception:
|
|
return "[retrieve] Invalid request format: must contain 'url2id', 'url_id', and 'goal'."
|
|
|
|
# Find every URL whose associated id equals the requested id
|
|
urls = [url for url, _id in url2id.items() if _id in target_id]
|
|
|
|
if not urls:
|
|
return f"[retrieve] id '{target_id}' not found in url2id."
|
|
|
|
# One URL → call readpage directly
|
|
# if len(urls) == 1:
|
|
# return self.readpage(urls[0], goal).strip()
|
|
|
|
|
|
# Several URLs share the same id → read them in parallel
|
|
url_page: List[dict] = []
|
|
non_existed_urls = []
|
|
url2page = None
|
|
if "url2page" in params:
|
|
url2page = params["url2page"]
|
|
for url in urls:
|
|
if url in url2page:
|
|
if isinstance(url2page[url], list):
|
|
page_content = ""
|
|
for item in url2page[url]:
|
|
page_content += str(item)
|
|
# page_content = "".join(url2page[url])
|
|
else:
|
|
page_content = str(url2page[url])
|
|
url_page.append({url: page_content})
|
|
else:
|
|
non_existed_urls.append(url)
|
|
else:
|
|
non_existed_urls = urls
|
|
|
|
with ThreadPoolExecutor(max_workers=3) as executor:
|
|
futures = {executor.submit(self.readpage, u, goal): u for u in non_existed_urls}
|
|
for future in as_completed(futures):
|
|
try:
|
|
url_page.append(future.result())
|
|
except Exception as e:
|
|
url_page.append(f"[retrieve] Error fetching {futures[future]}: {str(e)}")
|
|
|
|
output_str = ""
|
|
for item in url_page:
|
|
for key, value in item.items():
|
|
output_str += f"<id_{url2id[key]}>\nContent:\n{value}\n</id_{url2id[key]}>\n\n"
|
|
try:
|
|
if count_tokens(output_str) > max_tokens:
|
|
output_str = output_str[:max_tokens]
|
|
except Exception as e:
|
|
print(f"[retrieve] Error counting tokens: {str(e)}")
|
|
|
|
return output_str
|
|
|
|
def call_server(self, msgs, max_tries=10):
|
|
# 设置 OpenAI 的 API 密钥和 API 基础 URL 使用 vLLM 的 API 服务器。
|
|
openai_api_key = "EMPTY"
|
|
openai_api_base = "http://127.0.0.1:6002/v1"
|
|
client = OpenAI(
|
|
api_key=openai_api_key,
|
|
base_url=openai_api_base,
|
|
)
|
|
for attempt in range(max_tries):
|
|
try:
|
|
chat_response = client.chat.completions.create(
|
|
# TODO: [Important] If you change the summary model, you need to change the model path correspondingly.
|
|
model=os.getenv("SUMMARY_MODEL_PATH"),
|
|
messages=msgs,
|
|
temperature=0.7
|
|
)
|
|
content = chat_response.choices[0].message.content
|
|
if content:
|
|
try:
|
|
json.loads(content)
|
|
except:
|
|
# extract json from string
|
|
left = content.find('{')
|
|
right = content.rfind('}')
|
|
if left != -1 and right != -1 and left <= right:
|
|
content = content[left:right+1]
|
|
return content
|
|
except Exception as e:
|
|
# print(e)
|
|
if attempt == (max_tries - 1):
|
|
return ""
|
|
continue
|
|
|
|
def scraper_readpage(self, url: str) -> str:
|
|
max_retries = 3
|
|
|
|
scraper_reader = ScraperReader()
|
|
for attempt in range(max_retries):
|
|
page_content = scraper_reader.readpage(url)
|
|
if page_content:
|
|
return page_content
|
|
return "[visit] Failed to read page."
|
|
|
|
def readpage(self, url: str, goal=None, max_token=15000) -> str:
|
|
"""
|
|
Attempt to read webpage content by alternating between jina and aidata services.
|
|
|
|
Args:
|
|
url: The URL to read
|
|
goal: The goal/purpose of reading the page
|
|
|
|
Returns:
|
|
str: The webpage content or error message
|
|
"""
|
|
if isinstance(url, list):
|
|
url = url[0]
|
|
|
|
max_attempts = 10
|
|
for attempt in range(max_attempts):
|
|
content = self.scraper_readpage(url)
|
|
sevice = "scraper"
|
|
print(sevice, "length", len(content))
|
|
|
|
# Check if we got valid content
|
|
if content and not content.startswith("[visit] Failed to read page.") and content != "[visit] Empty content." and not content.startswith("[document_parser]"):
|
|
content = content[:int(os.getenv("WEBCONTENT_MAXLENGTH", 150000))]
|
|
|
|
return {url: content}
|
|
else:
|
|
return {url: ""}
|
|
|