# Created by Leandro Carneiro at 19/01/2024
# Description: 
# ------------------------------------------------
from langchain_openai import ChatOpenAI
from langchain_together import Together

import os

def invoke_llm(context, task, model):
    prompt = f"""You are an assistant of a newspaper. 
                 Execute the task just based on the given context.
                 The task is delimited by <> and the context is delimited by <>.
                 Write in a formal language and in portuguese language.
                 Execute the task just based on the given context.
                 Your task is: <{task}>
                 The context is: <{context}>
                 Answer here:
              """

    if model == 'openai':
        llm=ChatOpenAI(model_name="gpt-3.5-turbo",
                           temperature=0.3,
                           openai_api_key=os.environ['OPENAI_KEY'],
                           max_tokens=1000)
        result = llm.invoke(prompt)
        return result.content
    else:
        llm=Together(model="mistralai/Mixtral-8x7B-Instruct-v0.1",
                           temperature=0.3,
                           together_api_key=os.environ['TOGETHER_KEY'],
                           max_tokens=1000)
        result = llm.invoke(prompt)
        return result