File size: 604 Bytes
5ee5862
 
 
 
 
 
 
 
6eca139
 
5ee5862
 
6eca139
5ee5862
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# cleaner_chain.py
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) -> str:
        # Use invoke() instead of run() to comply with new LangChain practices
        return self.invoke({"kb_answer": kb, "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