import os
from langchain.chains import LLMChain
from langchain_groq import ChatGroq
from prompts import cleaner_prompt

class CleanerChain(LLMChain):
    def merge(self, kb: str, web: str, chat_history: list) -> str:
        """
        Merges the KB answer and web answer with chat history context.
        Appends the chat history before invoking the cleaner chain.
        """
        # Convert chat history into a formatted string
        history_context = "\n".join([f"User: {msg['content']}" for msg in chat_history])
        
        # Add the chat history context to the kb and web answers
        combined_input = f"{history_context}\nKB Answer: {kb}\nWeb Answer: {web}"
        
        # Use invoke() to merge the context with the knowledge base and web answers
        return self.invoke({"kb_answer": combined_input, "web_answer": web})

def get_cleaner_chain() -> CleanerChain:
    chat_groq_model = ChatGroq(model="Gemma2-9b-It", groq_api_key=os.environ["GROQ_API_KEY"])
    chain = CleanerChain(
        llm=chat_groq_model,
        prompt=cleaner_prompt
    )
    return chain