from typing import  Dict, List, Any
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace, HuggingFaceEndpoint
import torch
from huggingface_hub import login

api_key = 'hf'+ '_' + 'tlVzheuQBwjAxOtNKPqnHSQprFYnDLllut'
login(token=api_key)

class EndpointHandler:
    def __init__(self, path1="Cognute02/llama_3_1_8B_4bit"):
        # load model and processor from path
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype="float16",
            bnb_4bit_use_double_quant=True,
        )

        llm = HuggingFacePipeline.from_model_id(
                model_id=path1,
                task="text-generation",
                pipeline_kwargs=dict(
                    max_new_tokens=512,
                    do_sample=False,
                    repetition_penalty=1.03,
                    return_full_text=False,
                    temperature = 0.25
                ),
                model_kwargs={"quantization_config": quantization_config},
            )

        self.chatllm = ChatHuggingFace(llm=llm)

    def __call__(self, data):
        inputs = data['inputs']
        tools = data['tools']
        llm_ = self.chatllm.bind_tools(tools)
        outputs = llm_.invoke(inputs)
        return outputs