import psutil
from huggingface_hub import InferenceClient
import json
import streamlit as st
import os
from dotenv import load_dotenv


class QwenLlmHandler:

    def __init__(self):
        load_dotenv()
        self.client = InferenceClient(
            "Qwen/Qwen2.5-Coder-32B-Instruct",
            token=os.getenv("INFERENCE_API_TOKEN"),
        )
        st.info("Using LLM Qwen/Qwen2.5-Coder-32B-Instruct via inference API")

    def extract_data(self, text: str, entities: str):
        messages = [{
            "role": "user",
            "content": """Du bist ein NER-Model. Gebe die Veranstaltungsinformationen aus dem text in JSON Format zurück.
                            Es sollen keine Markdown Elemente enthalten sein, nur das JSON Objekt als string. 
                            Gebe Nur das JSON als Antwort zurück. JSON_Schema:  
                            {""" +
                       entities +
                       """"
                       }
                       Text: """ + text}]

        response = self.client.chat_completion(messages, max_tokens=1000)
        return json.loads(response.choices[0].message.content)