import os
from setup.environment import default_model
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI


async def aclaude_answer(claude_client, claude_context_model, prompt):
    print("\n\nComeçou uma requisição pelo Claude")
    response = await claude_client.messages.create(
        model=claude_context_model,
        max_tokens=100,  # Máximo é 4096
        messages=[{"role": "user", "content": prompt}],
    )
    return response.content[
        0
    ].text  # O response.content é uma lista pois é passada uma lista de mensagens, e também retornado uma lista de mensagens, sendo a primeira a mais recente, que é a resposta do model


async def agpt_answer(prompt):
    gpt = ChatOpenAI(
        temperature=0,
        model="gpt-4o-mini",
        api_key=os.environ.get("OPENAI_API_KEY"),
        max_retries=5,
    )
    response = await gpt.ainvoke([HumanMessage(content=prompt)])
    return response.content


def gpt_answer(
    prompt,
    temperature=0,
    model=default_model,
    max_retries=5,
    shouldReturnFullResponse=False,
):
    gpt = ChatOpenAI(
        temperature=temperature,
        model=model,
        api_key=os.environ.get("OPENAI_API_KEY"),
        max_retries=max_retries,
    )
    response = gpt.invoke([HumanMessage(content=prompt)])
    if shouldReturnFullResponse:
        return response
    else:
        return response.content