_utils/gerar_documento_utils/llm_calls.py

import os
from typing import Literal, cast

from pydantic import SecretStr
from _utils.langchain_utils.LLM_class import LLM, Google_llms
from setup.environment import default_model
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI

llm = LLM()


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):
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        gpt = ChatOpenAI(
            temperature=0,
            model="gpt-4o-mini",
            api_key=SecretStr(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,
):
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key:
        gpt = ChatOpenAI(
            temperature=temperature,
            model=model,
            api_key=SecretStr(api_key),
            max_retries=max_retries,
        )
        response = gpt.invoke([HumanMessage(content=prompt)])
        if shouldReturnFullResponse:
            return response
        else:
            return response.content


async def agemini_answer(
    prompt,
    model: Google_llms = "gemini-2.0-flash",
    temperature=0.4,
) -> str:
    gemini = llm.google_gemini(model, temperature)
    resposta = await gemini.ainvoke([HumanMessage(content=prompt)])

    if isinstance(resposta.content, list):
        resposta.content = "\n".join(resposta.content)  # type: ignore

    return cast(str, resposta.content)