Update LLM.py

LLM.py

@@ -2,102 +2,86 @@ 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_google_genai import ChatGoogleGenerativeAI
-import os
+from dotenv import load_dotenv
 from huggingface_hub import login
+import os
 
+# Load environment variables and authenticate
 load_dotenv()
-
-login(token=os.environ["HUGGING_FACE_API_KEY"])
+login(token=os.environ.get("HUGGING_FACE_API_KEY", ""))
 os.environ['CURL_CA_BUNDLE'] = ''
 
-load_dotenv()
-
-class Bot():
+class Bot:
     def __init__(self):
-        self.groq_models = ['gemma-7b-it', 'llama3-70b-8192',\
-                            'llama3-8b-8192', 'mixtral-8x22b']
-        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",\
-                              "gemini-1.5-pro"]
-        self.models = [
-    "gemini-pro", "gemini-pro-vision", 
-    "01-ai/Yi-1.5-34B-Chat", "google/gemma-1.1-2b-it", 
-    "gemma-7b-it", "llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768"]    
-
-
-    def call_groq(self, model, temp = 0.7, given_prompt = "Hi"):
+        # Updated model lists: Remove gated or unsupported IDs        .groq_models = [
+            'gemma-7b-it',
+            'llama3-70b-8192',
+            'llama3-8b-8192',
+            'mixtral-8x22b'
+        ]
+        self.hf_models = [
+            "01-ai/Yi-1.5-34B-Chat",
+            "google/gemma-1.1-2b-it"
+        ]
+        # Use supported Google GenAI model names
+        self.google_models = [
+            "gemini-pro",
+            "gemini-pro-vision"
+        ]
+        # Master list for sampling (only include accessible models)
+        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(
-                temperature=temp,
-                model= model
-            )
-
-            system = "You are a helpful assistant."
-            human = "{text}"
-            prompt = ChatPromptTemplate.from_messages([("system", system), ("human", human)])
-
+            llm = ChatGroq(model=model, temperature=temp)
+            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"⚠️ [Groq:{model}] {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
-                )
-
+            llm = HuggingFaceEndpoint(repo_id=model, temperature=temp)
             chat = ChatHuggingFace(llm=llm, verbose=True)
-
             template = """
-                You are a helpful assistant
-
-                User: {query}
-
-                Answer:
-            """
+You are a helpful assistant.
 
-            prompt = PromptTemplate(
-                template=template,
-                input_variables=["query"]
-            )
-
-            chain =prompt | chat | StrOutputParser()
+User: {query}
 
+Answer:
+"""
+            prompt = PromptTemplate(template=template, input_variables=["query"])
+            chain = prompt | chat | StrOutputParser()
             return chain.invoke({"query": given_prompt})
-
         except Exception as e:
-            return f"Error: {str(e)}"
+            return f"⚠️ [HF:{model}] {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()
+            gm = ChatGoogleGenerativeAI(model=model, temperature=temp)
+            prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful assistant."), ("human", "{text}")])
+            chain = prompt | gm | StrOutputParser()
             return chain.invoke({"text": given_prompt})
         except Exception as e:
-            return f"Error: {str(e)}"
-
-    def response(self, model, prompt="Hi", temprature = 0.7):
-        if model in self.groq_models:
-            res_show = self.call_groq(temp = temprature, given_prompt = prompt, model= model)
-        elif model in self.hf_models:
-            res_show = self.call_hf(given_prompt = prompt, temp = temprature, model = model)
-        elif model in self.google_models:
-            res_show = self.call_google(given_prompt = prompt, temp = temprature, model = model)
-        else:
-            return "Sorry! App not working properly"
-        return res_show
+            return f"⚠️ [Google:{model}] {str(e)}"
 
+    def response(self, model, prompt="Hi", temperature=0.7):
+        # Route to the correct provider and catch errors
+        try:
+            if model in self.groq_models:
+                return self.call_groq(model, temp=temperature, given_prompt=prompt)
+            if model in self.hf_models:
+                return self.call_hf(model, temp=temperature, given_prompt=prompt)
+            if model in self.google_models:
+                return self.call_google(model, temp=temperature, given_prompt=prompt)
+            return f"❌ Unsupported model: {model}"
+        except Exception as e:
+            return f"⚠️ Skipping `{model}` due to error: {str(e)}"