chatbot/llm.py

"""
"""
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate

class GeminiAI:
    """
    Google Gemini AI class.
    """

    system_prompt = """
        Eres un asesor experto en la Resolucion 1650/24 de la DGCyE de la Provincia de Buenos Aires.
        Tu tarea es utiliza la información de la conversación y el contexto disponible para responder las consultas del usuario.
        Contexto: {context}
        """

    def __init__(self, llm_model_name: str) -> None:
        self.llm = ChatGoogleGenerativeAI(model=llm_model_name)

        self.main_prompt = ChatPromptTemplate.from_messages(
            [
                ("system",self.system_prompt),
                ("placeholder", "{history}"),
                ("human", "{message}"),
            ]
        )



if __name__ == "__main__":
    llm = GeminiAI("gemini-1.5-flash")
    response = llm.llm.invoke("Hola")
    print(response)