Update LLM.py

LLM.py

@@ -1,101 +1,92 @@
 from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
 from langchain_groq import ChatGroq
-from langchain_huggingface import ChatHuggingFace
-from langchain_huggingface import HuggingFaceEndpoint
-from dotenv import load_dotenv
-from langchain.schema.output_parser import StrOutputParser
-from langchain_huggingface import ChatHuggingFace
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from langchain_google_genai import ChatGoogleGenerativeAI
-import os
+from langchain.schema.output_parser import StrOutputParser
+from dotenv import load_dotenv
 from huggingface_hub import login
+import os
 
+# Load environment variables
 load_dotenv()
-
 login(token=os.environ["HUGGING_FACE_API_KEY"])
-os.environ['CURL_CA_BUNDLE'] = ''
-
-load_dotenv()
+os.environ['CURL_CA_BUNDLE'] = ''  # Optional workaround for certificate issues
 
 class Bot():
     def __init__(self):
-        self.groq_models = ['gemma-7b-it', 'llama3-70b-8192',\
-                            'llama3-8b-8192', 'mixtral-8x7b-32768']
-        self.hf_models = ["01-ai/Yi-1.5-34B-Chat", "google/gemma-1.1-2b-it",\
+        self.groq_models = ['gemma-7b-it', 'llama3-70b-8192',
+                            'llama3-8b-8192', 'mixtral-8x22b']  # Updated here
+        self.hf_models = ["01-ai/Yi-1.5-34B-Chat", "google/gemma-1.1-2b-it",
                           "google/gemma-1.1-7b-it"]
-        self.google_models = ["gemini-1.0-pro", "gemini-1.5-flash",\
+        self.google_models = ["gemini-1.0-pro", "gemini-1.5-flash",
                               "gemini-1.5-pro"]
-        self.models = ["gemini-1.0-pro", "gemini-1.5-flash", "gemini-1.5-pro", "01-ai/Yi-1.5-34B-Chat", "google/gemma-1.1-2b-it",\
-                          "google/gemma-1.1-7b-it", 'gemma-7b-it', 'llama3-70b-8192', 'llama3-8b-8192', 'mixtral-8x7b-32768']
 
-    def call_groq(self, model, temp = 0.7, given_prompt = "Hi"):
+        # Main model pool
+        self.models = self.google_models + self.hf_models + self.groq_models
+
+    def call_groq(self, model, temp=0.7, given_prompt="Hi"):
         try:
             llm = ChatGroq(
+                model=model,
                 temperature=temp,
-                model= model
+                groq_api_key=os.environ["GROQ_API_KEY"]
             )
-
-            system = "You are a helpful assistant."
-            human = "{text}"
-            prompt = ChatPromptTemplate.from_messages([("system", system), ("human", human)])
-
+            prompt = ChatPromptTemplate.from_messages([
+                ("system", "You are a helpful assistant."),
+                ("human", "{text}")
+            ])
             chain = prompt | llm | StrOutputParser()
             return chain.invoke({"text": given_prompt})
 
         except Exception as e:
-            return f"Error: {str(e)}"
+            return f"Error (Groq): {str(e)}"
 
-    def call_hf(self,model, temp = 0.7, given_prompt = "Hi"):
+    def call_hf(self, model, temp=0.7, given_prompt="Hi"):
         try:
             llm = HuggingFaceEndpoint(
                 repo_id=model,
                 temperature=temp
-                )
-
+            )
             chat = ChatHuggingFace(llm=llm, verbose=True)
+            prompt = PromptTemplate(
+                template="""
+                    You are a helpful assistant
 
-            template = """
-                You are a helpful assistant
-
-                User: {query}
-
-                Answer:
-            """
+                    User: {query}
 
-            prompt = PromptTemplate(
-                template=template,
+                    Answer:
+                """,
                 input_variables=["query"]
             )
-
-            chain =prompt | chat | StrOutputParser()
-
+            chain = prompt | chat | StrOutputParser()
             return chain.invoke({"query": given_prompt})
 
         except Exception as e:
-            return f"Error: {str(e)}"
+            return f"Error (HF): {str(e)}"
 
-    def call_google(self,model, temp=0.7, given_prompt = "Hi"):
+    def call_google(self, model, temp=0.7, given_prompt="Hi"):
         try:
-            model = ChatGoogleGenerativeAI(model = model, temprature = temp)
-            system = "You are a helpful assistant."
-            human = "{text}"
-            prompt = ChatPromptTemplate.from_messages([("human", human)])
-            chain = prompt | model | StrOutputParser()
+            llm = ChatGoogleGenerativeAI(
+                model=model,
+                temperature=temp,
+                google_api_key=os.environ["GOOGLE_API_KEY"]
+            )
+            prompt = ChatPromptTemplate.from_messages([
+                ("human", "{text}")
+            ])
+            chain = prompt | llm | StrOutputParser()
             return chain.invoke({"text": given_prompt})
+
         except Exception as e:
-            return f"Error: {str(e)}"
+            return f"Error (Google): {str(e)}"
 
-    def response(self, model, prompt="Hi", temprature = 0.7):
+    def response(self, model, prompt="Hi", temperature=0.7):
         if model in self.groq_models:
-            res_show = self.call_groq(temp = temprature, given_prompt = prompt, model= model)
+            return self.call_groq(model=model, temp=temperature, given_prompt=prompt)
         elif model in self.hf_models:
-            res_show = self.call_hf(given_prompt = prompt, temp = temprature, model = model)
+            return self.call_hf(model=model, temp=temperature, given_prompt=prompt)
         elif model in self.google_models:
-            res_show = self.call_google(given_prompt = prompt, temp = temprature, model = model)
+            return self.call_google(model=model, temp=temperature, given_prompt=prompt)
         else:
-            return "Sorry! App not working properly"
-        return res_show
-
-
-
-
+            return "Sorry! App not working properly – unknown model"