# -*- coding: utf-8 -*-
# entity_recognition.py
# @author: lyg
# @date: 2025-04-24
# @version: 0.0.1
# @description: 实体抽取，将文本中的实体进行识别和提取。

from langchain_openai.chat_models import ChatOpenAI
from langchain_core.prompts import HumanMessagePromptTemplate, ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
import json

from knowledgebase import utils


class EntityRecognition:
    """
    实体识别抽取。

    使用langchain构建实体抽取流程。
    """
    cache_file = "entity_recognition.cache"

    def __init__(self):
        llm = ChatOpenAI(temperature=0,
                         model="qwen2.5-72b-instruct",
                         base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                         api_key="sk-15ecf7e273ad4b729c7f7f42b542749e")
        msg = HumanMessagePromptTemplate.from_template(template="""
# 指令
请从给定的文本中提取实体词列表。
# 约束
- 输出格式为JSON格式；
- 输出数据结构为字符串数组。
# 示例
```json
["实体1","实体2"]
```

# 文本如下：
{text}
"""
                                                       )
        prompt = ChatPromptTemplate.from_messages([msg])
        parser = JsonOutputParser(pydantic_object=list[str])
        self.chain = prompt | llm | parser
        self.cache = {}
        self.load_cache()

    def load_cache(self):
        """
        加载缓存。
        """
        if utils.file_exists(self.cache_file):
            text = utils.read_from_file(self.cache_file)
            self.cache = json.loads(text)

    def save_cache(self):
        """
        保存缓存。
        """
        text = json.dumps(self.cache)
        utils.save_to_file(text, self.cache_file)

    def run(self, in_text: str) -> list[str]:
        """
        运行实体识别抽取。
        """
        # 缓存命中
        text_md5 = utils.generate_md5(in_text)
        if text_md5 in self.cache:
            return self.cache[text_md5]
        result = self.chain.invoke({"text": in_text})
        self.cache[text_md5] = result
        self.save_cache()
        return result