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

import config, debug, prompt

DEBUG = True

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

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

def get_task() -> str:
    # 从引擎中获取任务的接口
    # 返回的类型应为字符串或None
    raise NotImplementedError

def push_result(result):
    # 将提取结果返回给引擎的接口
    raise NotImplementedError

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 response

def run():
    if DEBUG:
        input_text = debug.input_text
    else:
        input_text = get_task()
        if not input_text:
            return
    response = entity_extract(input_text)
    if DEBUG:
        print(response)
    else:
        push_result(response)

if __name__ == "__main__":
    if DEBUG:
        run()
        exit()
    while True:
        run()
        time.sleep(5)