File size: 1,739 Bytes
80a598c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import os
import re
from typing import List
import dashscope
def rewrite_query(query: str, n_max_query: int = 5) -> List[str]:
raw_prompt = _build_raw_prompt(query)
try:
dashscope.base_http_api_url = 'https://poc-dashscope.aliyuncs.com/api/v1'
response_query = dashscope.Generation.call(
api_key='sk-6bddfc116de744c3aa1d66893cc87b20',
model='pre-qwen-rag-rewrite-decomp-chat',
prompt=raw_prompt,
use_raw_prompt=True
)
queries = _parse_output(response_query.output.text)
return queries[:n_max_query]
except Exception as e:
raise ValueError(f'Call dashscope failed: {str(e)}')
def _build_raw_prompt(query: str) -> str:
raw_prompt = '<|im_start|>system\n# Tool\n[\n {\n \"type\": \"function\",\n \"function\": {\n \"name\": \"search\",\n \"description\": \"Utilize the web search engine to retrieve relevant information based on multiple queries\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"queries\": {\n \"type\": \"array\",\n \"items\": {\n \"type\": \"string\",\n \"description\": \"A search query\"\n },\n \"description\": \"The list of search queries\"\n }\n },\n \"required\": [\n \"queries\"\n ]\n }\n }\n }\n]\n<|im_end|>\n<|im_start|>user\n'+query+'<|im_end|>\n<|im_start|>assistant\n# Tool Call\n'
return raw_prompt
def _parse_output(output: str) -> List[str]:
results = re.findall('\'(.*?)\'', output)
return results
if __name__ == '__main__':
print(rewrite_query('深度分析kimi chat的发展现状'))
|