from langchain.chains import RetrievalQA
from langchain_groq import ChatGroq
from app.config import config
from modules.processing.llm1.config import GROQ_API_KEY
from modules.processing.llm1.prompt_manager import PromptManager

class RetrievalQAProcessor:
    """Ejecuta consultas en los tweets utilizando FAISS y Groq."""

    def __init__(self, vectorstore, cantidad_categorias):
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=ChatGroq(
                groq_api_key=GROQ_API_KEY,
                model_name=config.GROQ_MODEL_LLM1
            ),
            chain_type="stuff",
            retriever=vectorstore.as_retriever(),
            return_source_documents=False,
            chain_type_kwargs={"prompt": PromptManager.get_prompt(cantidad_categorias)}
        )

    def query(self, pregunta: str) -> str:
        respuesta = self.qa_chain.invoke(pregunta)
        return respuesta['result']