from dotenv import load_dotenv

load_dotenv()
import time
from openai import OpenAI
from datetime import datetime

import config, debug, prompt, json

"""
备注：
    目前使用ollama部署的qwen2.5:32b-instruct-q5_K_M, 
    如果输入字符串过长(大致超过1000个字符或1500token)会因上下文窗口过短导致提示词被遗忘，
    进而输出不合规定的结果！
"""

llm = OpenAI(base_url=config.LLM_BASE_URL, api_key=config.API_KEY)


class entityExtractionProcess:
    def entity_extract(input_text):
        messages = [
            {"role": "system", "content": prompt.ENTITY_EXTRACT},
            {"role": "system", "content": f"今天的日期是：{str(datetime.today())}"},
            {"role": "user", "content": f"{input_text}"},
        ]

        timenow = time.time()

        response = (
            llm.chat.completions.create(
                model=config.MODEL, messages=messages, temperature=0, max_tokens=128_000
            )
            .choices[0]
            .message.content
        )
        print(f"本次输出花费时间：{time.time() - timenow} 秒")

        return json.loads(response)