{"repo": "casanovalll/lawagent", "files": {"LawAgent/Tools/recursive_classifier.py": {"3e2b5b3f6fe88516e5dd0f47f77f7c9fb438e2f5": {"t": 1786068483.2525494, "n": 6663, "h": "fd76b73db020c1fd", "c": "import json\r\nimport os\r\n\r\nimport json5\r\nimport requests\r\nimport editdistance\r\nfrom typing import Union\r\n\r\nfrom qwen_agent.tools.base import BaseTool, register_tool\r\nfrom openai import OpenAI, OpenAIError\r\n\r\nfrom LawAgent.Utils import build_username\r\nfrom tqdm import tqdm\r\n\r\n\r\n@register_tool('cause_classifier')\r\nclass RecursiveClassifier(BaseTool):\r\n    name = 'cause_classifier'\r\n    description = '一个案由分类器，把一个事件分到某一个案由'\r\n    parameters = [{\r\n        'name': 'query',\r\n        'type': 'string',\r\n        'description':\r\n            '需要分类的文本',\r\n        'required': True\r\n    }]\r\n\r\n    def __init__(self, tree_path: str, *args, **kwargs):\r\n        super().__init__(*args, **kwargs)\r\n        self.deepseek_key = os.environ.get('DEEPSEEK_API_KEY', kwargs.get('deepseek_api_key', None))\r\n        self.base_url = os.environ.get('DEEPSEEK_BASE_URL', kwargs.get('deepseek_base_url', \"https://api.deepseek.com\"))\r\n        with open(tree_path, 'r', encoding='utf-8') as f:\r\n            self.tree = json5.load(f)\r\n\r\n    def call(self, params: Union[str, dict], **kwargs) -> Union[str, list, dict]:\r\n        if isinstance(params, str):\r\n            params = json5.loads(params)\r\n        prompt = params['query']\r\n        return self.classify(prompt, self.tree[\"root\"])\r\n\r\n    def classify(self, prompt: str, node: dict) -> str:\r\n        def _find_nearest_node(node_list: list, query):\r\n            min_distance = 1\r\n            _result = None\r\n            for _ in node_list:\r\n                distance = editdistance.eval(query, _) / max(len(query), len(_))\r\n                if distance < min_distance:\r\n                    min_distance = distance\r\n                    _result = _\r\n            return _result\r\n\r\n        # print(\"Classifier prompt\", prompt, node)\r\n        description = node['description']\r\n        next_nodes = node['next']\r\n        if len(next_nodes) == 1:\r\n            return next_nodes[0]\r\n        for _ in range(3):\r\n            try:\r\n                client = OpenAI(\r\n                    api_key=self.deepseek_key,\r\n                    base_url=self.base_url,\r\n                )\r\n\r\n                is_multy = node.get('is_multy', False)\r\n                base_prompt = BASE_PROMPT if not is_multy else MULTI_PROMPT\r\n                sys_prompt = base_prompt.format(description=description, choice=next_nodes)\r\n                sys_prompt += BASE_OUTPUT_EXAMPLE if not is_multy else MULTY_OUTPUT_EXAMPLE\r\n                response = client.chat.completions.create(\r\n                    model=\"deepseek-chat\",\r\n                    messages=[\r\n                        {\"role\": \"system\", \"content\": sys_prompt},\r\n                        {\"role\": \"user\", \"content\": prompt}\r\n                    ],\r\n                    temperature=0.1,\r\n                    max_tokens=512,\r\n                    stream=False,\r\n                    response_format={\r\n                        'type': 'json_object'\r\n                    }\r\n                )\r\n                next_node = json5.loads(response.choices[0].message.content)['choice']\r\n                print(\"Classifier next node\", next_node)\r\n                if is_multy:\r\n                    assert isinstance(next_node, list)\r\n                    next_node = [_find_nearest_node(next_nodes, _) for _ in next_node]\r\n                    return next_node\r\n                assert isinstance(next_node, str)\r\n                next_node = _find_nearest_node(next_nodes, next_node)\r\n                if next_node in node['next']:\r\n                    if next_data := self.tree.get(next_node, None):\r\n                        return self.classify(prompt, next_data)\r\n                    return next_node\r\n\r\n            except (OpenAIError, AssertionError):\r\n                pass\r\n        return next_nodes[0]\r\n\r\n\r\nBASE_PROMPT = r\"\"\"\r\n你是一个案件分类器，你需要根据提供的事件文本和背景知识，把事件分类到待选案由中的一个。\r\n请输出一个json格式的字符串\r\n背景知识：\r\n{description}\r\n待选案由：\r\n{choice}\r\n\r\n\"\"\"\r\n\r\nMULTI_PROMPT = r\"\"\"\r\n你是一个案件分类器，你需要根据提供的事件文本和背景知识，把事件分类到待选案由中的一个或多个。\r\n请输出一个json格式的字符串\r\n背景知识：\r\n{description}\r\n待选案由：\r\n{choice}\r\n\r\n\"\"\"\r\nBASE_OUTPUT_EXAMPLE = \"\"\"\r\n输出样例：\r\n{\r\n    \"choice\":\"垄断纠纷\"\r\n}\r\n\"\"\"\r\nMULTY_OUTPUT_EXAMPLE = \"\"\"\r\n输出样例：\r\n{\r\n    \"choice\":[\"垄断纠纷\"]\r\n}\r\n\"\"\"\r\n\r\n\r\ndef _description_generate(flatten_tree: dict, reference_dir: str = None):\r\n    from zhipuai import ZhipuAI\r\n    system_prompt = \"\"\"\r\n    利用搜索引擎，为将一个案件分到以下的分类编写背景知识，主要包括类别的简介和类别间的差异\r\n    - 用精炼简洁的语言\r\n    \"\"\"\r\n    references = os.listdir(reference_dir) if reference_dir else None\r\n\r\n    for k, v in tqdm(flatten_tree.items(), total=len(flatten_tree)):\r\n        try:\r\n            if description := v['description']:\r\n                if references and description in references:\r\n                    with open(os.path.join(reference_dir, description), 'r', encoding='utf-8') as f:\r\n                        description = f.read()\r\n                        v['description'] = description\r\n                continue\r\n            if not v[\"next\"] or len(v['next']) <= 1:\r\n                continue\r\n            query = f\"为以下{k}的子案由编写背景知识，请用精炼简洁的语言，同时注意区分不同类别之间的差异：\\n{v['next']}\"\r\n            print(query)\r\n            client = ZhipuAI(api_key=os.environ[\"ZHIPU_APIKEY\"])\r\n            response = client.chat.completions.create(\r\n                model=\"glm-4-0520\",  # 填写需要调用的模型名称\r\n                messages=[\r\n                    {\"role\": \"system\",\r\n                     \"content\": system_prompt},\r\n                    {\"role\": \"user\",\r\n                     \"content\": query},\r\n                ],\r\n                tools=[{\"type\": \"web_search\", \"web_search\": {\"search_result\": True}}],\r\n                stream=True,\r\n                top_p=0.7,\r\n                temperature=0.95,\r\n                max_tokens=1024,\r\n            )\r\n            answer_parts = []\r\n            print(\"--------------------------------------------------------------------------------\")\r\n            for chunk in response:\r\n                if content := chunk.choices[0].delta.content:\r\n                    print(content, end=\"\")\r\n                    answer_parts.append(content)\r\n            response_str = \"\".join(answer_parts)\r\n            print(\"\")\r\n            print(k)\r\n            print(\"--------------------------------------------------------------------------------\")\r\n            v['description'] = response_str\r\n        except Exception as e:\r\n            print(e)\r\n            print(\"error\", k)\r\n\r\n    return flatten_tree\r\n\r\n\r\nif __name__ == '__main__':\r\n    with open(\"分类树-细.json\", \"r\") as f:\r\n        flatten_tree = json5.load(f)\r\n    flatten_tree = _description_generate(flatten_tree, reference_dir=\"reference\")\r\n    with open(\"flatten_with_desc.json\", \"w\") as f:\r\n        f.write(json.dumps(flatten_tree, ensure_ascii=False, indent=4))\r\n"}}}}