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