| 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的发展现状')) | |