import os
from langchain_openai import ChatOpenAI
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_together import Together
from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings

from dotenv import load_dotenv

load_dotenv()

class PersonalModels:
    def get_chat_model(self, 
                       model: str = 'mistralai/Mistral-7B-Instruct-v0.2', 
                       api_key: str = os.environ.get('TOGETHER_API_KEY', ''), 
                       max_tokens: int = 1000, 
                       temperature: float = 0.0, 
                       verbose: bool = True) -> ChatOpenAI:
        llm = ChatOpenAI(
                        api_key=api_key,
                        base_url="https://api.together.xyz/v1",
                        model=model,
                        max_tokens=max_tokens,
                        temperature=temperature,
                        verbose=verbose
                        )
        return llm

    def get_embedding_model(self, 
                            model: str = 'intfloat/multilingual-e5-large') -> HuggingFaceInferenceAPIEmbeddings:
        embeddings = HuggingFaceInferenceAPIEmbeddings(api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN', ''), model_name=model)
        # embeddings = HuggingFaceEmbeddings(model_name=model,model_kwargs = {'device': 'cuda'})
        # print('modelo está na gpu?',next(embeddings.model.parameters()).is_cuda)
        return embeddings

    def get_rag_model(self, 
                      model: str = 'mistralai/Mistral-7B-Instruct-v0.2', 
                      api_key: str = os.environ.get('TOGETHER_API_KEY', ''), 
                      max_tokens: int = 1000, 
                      temperature: float = 0.0, 
                      top_k: int = 20, 
                      verbose: bool = True) -> Together:
        llm = Together(
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        top_k=top_k,
        together_api_key=api_key,
        verbose=verbose
        )
        return llm