TypeChat让AI输出结构化数据

typechat.py

import json
import logging
import re

import js2py

LOG = logging.getLogger(__name__)

RE_CODE_BLOCK = re.compile(
    r'```(json)?\n?(.+)\n?```', re.MULTILINE | re.DOTALL
)


class TypeChat:
    """结构化的对话提示"""

    def system_prompt(self, typescript: str):
        template = """
You are a service that translates user requests into JSON object according to the following TypeScript definitions of type `Result`.
你是一个用户请求翻译器，根据用户请求和以下TypeScript定义准确翻译成符合 `Result` 类型的JSON对象。TypeScript定义如下：
```
{TYPESCRIPT}
```
Please translate the user request into a JSON object directly, with 2 spaces of indentation, use `null` to represent values that cannot be translated. Please make sure response JSON text directly.
请将用户请求直接翻译成JSON数据，用2个空格缩进，无法翻译的值用null表示。请确保直接输出JSON数据，不要回复其他内容。
"""
        return template.strip().format(TYPESCRIPT=typescript.strip())

    def user_prompt(self, user_request: str):
        template = """
Please translate the following user request into JSON object（请翻译以下用户请求成JSON数据）:
{USERREQUEST}
"""
        return template.strip().format(USERREQUEST=user_request.strip())

    def _extract_markdown_code_block(self, response: str):
        """
        用正则表达式提取 ```json``` 之间的内容
        """
        match = RE_CODE_BLOCK.search(response)
        if match:
            return match.group(2).strip()
        return response

    def _load_json(self, text: str):
        result = error = None
        json_text = text.replace('undefined', 'null')
        try:
            result = json.loads(json_text)
        except (json.JSONDecodeError, Exception) as ex:
            error = ex
            LOG.info('json failed to parse response: %s', text)
        if result is not None:
            return None, result
        js_text = f'_result={text}'
        try:
            result = js2py.eval_js(js_text)
            result = result.to_dict()
        except (js2py.PyJsException, Exception) as ex:
            error = ex
            LOG.info('js2py failed to parse response: %s', text)
        if result is not None:
            return None, result
        return error, None

    def parse_response(self, response: str):
        """
        提取回复数据，返回 (error, result)
        """
        text = self._extract_markdown_code_block(response)
        error, result = self._load_json(text)
        if error is not None:
            return error, None
        if 'Result' in result:
            result = result['Result']
        return None, result