Update chatbot.py

chatbot.py

@@ -1,28 +1,25 @@
-import warnings
-warnings.filterwarnings("ignore", category=DeprecationWarning)
 import os
+import time
+import json
 from groq import Groq
 from langchain.memory import ConversationTokenBufferMemory
-from langchain.memory import ConversationBufferMemory
 from langchain_community.chat_models import ChatOpenAI
-from langdetect import detect
-from deep_translator import GoogleTranslator
 from langchain_community.document_loaders.csv_loader import CSVLoader
 from langchain_community.vectorstores import FAISS
-import time
-import json
+from deep_translator import GoogleTranslator
+
 
 class Comsatsbot:
     def __init__(self, hf, llm, api_keys, chats_collection, paths, index_path='faiss_kb'):
         self.llm = llm
         self.api_keys = api_keys
         self.client = None
-        self.models = ["llama3-groq-70b-8192-tool-use-preview", "llama-3.1-70b-versatile", "llama3-70b-8192"]
-
-
-        # Initialize memory buffer and MongoDB connection
-        self.memory = ConversationTokenBufferMemory(llm=self.llm, max_token_limit=3000 )
-        #self.memory = ConversationBufferMemory( return_messages=True, memory_key="chat_history")
+        self.models = [
+            "llama3-groq-70b-8192-tool-use-preview",
+            "llama-3.1-70b-versatile",
+            "llama3-70b-8192"
+        ]
+        self.memory = ConversationTokenBufferMemory(llm=self.llm, max_token_limit=3000)
         self.chats_collection = chats_collection
         self.index_path = index_path
         self.hf = hf
@@ -42,19 +39,16 @@ class Comsatsbot:
     def initialize_faiss_index(self):
         if os.path.exists(self.index_path):
             self.faiss_index = FAISS.load_local(self.index_path, self.hf, allow_dangerous_deserialization=True)
-            self.faiss_retriever = self.faiss_index.as_retriever(search_kwargs={"k": 5})
         else:
             documents = self.load_data(self.paths)
             self.faiss_index = FAISS.from_documents(documents, self.hf)
             self.faiss_index.save_local(self.index_path)
-            self.faiss_retriever = self.faiss_index.as_retriever(search_kwargs={"k": 5})
+        self.faiss_retriever = self.faiss_index.as_retriever(search_kwargs={"k": 5})
 
     def retrieve_answer(self, query):
         if self.faiss_retriever:
             return self.faiss_retriever.invoke(query)
-        else:
-            print("FAISS retriever is not initialized. Please create or load an index.")
-            return None
+        return None
 
     def create_chat_record(self, chat_id):
         self.chats_collection.insert_one({
@@ -75,20 +69,14 @@ class Comsatsbot:
         return chat_record.get('history', [])
 
     def new_chat(self, chat_id):
-        # Check if chat ID already exists
         if self.chats_collection.find_one({"_id": chat_id}):
-            raise KeyError(f"Chat ID {chat_id} exist already.")
-
-        # Create a new chat record if it doesn't exist
+            raise KeyError(f"Chat ID {chat_id} exists already.")
         self.create_chat_record(chat_id)
         return "success"
 
     def delete_chat(self, chat_id):
-        # Check if the chat ID exists
         if not self.chats_collection.find_one({"_id": chat_id}):
             raise KeyError(f"Chat ID {chat_id} does not exist.")
-
-        # Delete the chat if it exists
         self.chats_collection.delete_one({"_id": chat_id})
         return "success"
 
@@ -133,141 +121,107 @@ class Comsatsbot:
                 '''
         while True:
             for api_key in self.api_keys:
-                self.client = Groq(api_key=api_key)
-
-                for model in self.models:
-                    try:
-                        chat_completion = self.client.chat.completions.create(
+            self.client = Groq(api_key=api_key)
+            for model in self.models:
+                try:
+                    chat_completion = self.client.chat.completions.create(
                         messages=[
-                              {
-                    "role": "system",
-                    "content": prompt,
-                },
-                {
-                    "role": "user",
-                    "content": f"Answer the following question : {question}",
-                },
+                            {"role": "system", "content": prompt},
+                            {"role": "user", "content": f"Answer the following question: {question}"}
                         ],
                         model=model,
                         max_tokens=1024,
                     )
-
-                        response_content = chat_completion.choices[0].message.content
-                        return response_content
-                    except Exception as e:
-                        time.sleep(2)
-
-                        continue
-
-            return "Sorry, unable to provide an answer at this time."
+                    return chat_completion.choices[0].message.content
+                except Exception:
+                    time.sleep(2)
+                    continue
+        return "Sorry, unable to provide an answer at this time."
 
     def detect_language(self, question):
-
-        while True:
-            for api_key in self.api_keys:
-                self.client = Groq(api_key=api_key)
-    
-                for model in self.models:
-                    try:
-                        chat_completion = self.client.chat.completions.create(
+        for api_key in self.api_keys:
+            self.client = Groq(api_key=api_key)
+            for model in self.models:
+                try:
+                    chat_completion = self.client.chat.completions.create(
                         messages=[
                             {
-                    "role": "system",
-                    "content": """ 
-                    You are an expert agent and your task is to detect the language from the following provided question. Your generated output should be according to following json format.
-                    Json Output Format:
-                    {'detected_language': 'write the just detected language name: urdu, english '}
-
-
-                    
-                    """,
-                },
-                {
-                    "role": "user",
-                    "content": f"detect the language for the following Question: {question}",
-                },
+                                "role": "system",
+                                "content": """
+You are an expert agent and your task is to detect the language. 
+Return a JSON: {'detected_language': 'urdu' or 'english'}
+"""
+                            },
+                            {
+                                "role": "user",
+                                "content": f"Detect the language for: {question}"
+                            }
                         ],
                         model=model,
-                        max_tokens=512,
+                        max_tokens=256,
                         response_format={"type": "json_object"},
                     )
-    
-                        response_content = chat_completion.choices[0].message.content
-                        json_response = json.loads(response_content)
-                        return json_response['detected_language'].lower()
-                    except Exception as e:
-                        time.sleep(2)
-    
-                        continue
-    def translate_urdu(self, text):
+                    response = json.loads(chat_completion.choices[0].message.content)
+                    return response['detected_language'].lower()
+                except Exception:
+                    time.sleep(2)
+                    continue
+        return "english"
 
-        while True:
-            for api_key in self.api_keys:
-                self.client = Groq(api_key=api_key)
-    
-                for model in self.models:
-                    try:
-                        chat_completion = self.client.chat.completions.create(
+    def translate_urdu(self, text):
+        for api_key in self.api_keys:
+            self.client = Groq(api_key=api_key)
+            for model in self.models:
+                try:
+                    chat_completion = self.client.chat.completions.create(
                         messages=[
                             {
-                    "role": "system",
-                    "content": """ 
-                    You are an expert agent and your task is to translayte the following text in in proper urdu text and i want corrected urdu and dont need to add irrelevant text. Your generated output should be according to following json format.
-                    Json Output Format:
-                    {'text': 'write the just translated urdu text here and dont need to add irrelevant text'}
-
-
-                    
-                    """,
-                },
-                {
-                    "role": "user",
-                    "content": f"detect the language for the following Question: {text}",
-                },
+                                "role": "system",
+                                "content": """
+Translate the following text into proper Urdu. Return a JSON: 
+{'text': 'translated urdu text'}
+"""
+                            },
+                            {
+                                "role": "user",
+                                "content": f"Translate this: {text}"
+                            }
                         ],
                         model=model,
                         max_tokens=512,
                         response_format={"type": "json_object"},
                     )
-    
-                        response_content = chat_completion.choices[0].message.content
-                        json_response = json.loads(response_content)
-                        return json_response['text']
-                    except Exception as e:
-                        time.sleep(2)
-    
-                        continue
-
+                    response = json.loads(chat_completion.choices[0].message.content)
+                    return response['text']
+                except Exception:
+                    time.sleep(2)
+                    continue
+        return text
 
     def response(self, question, chat_id):
         chat_history = self.load_chat(chat_id)
 
-        # Load the previous conversation into memory
         for entry in chat_history:
             self.memory.save_context({"input": entry["question"]}, {"output": entry["answer"]})
 
         language = self.detect_language(question)
+
         if language == 'urdu':
             question_translation = GoogleTranslator(source='ur', target='en').translate(question)
             context = self.faiss_retriever.invoke(question_translation)
         else:
             context = self.faiss_retriever.invoke(question)
 
-        all_content = ''
-        for document in context:
-            page_content = document.page_content
-            all_content += page_content + '\n' 
-        answer = self.generate_response(question, self.memory.load_memory_variables({})['history'], all_content)
+        combined_context = '\n'.join([doc.page_content for doc in context])
+        answer = self.generate_response(question, self.memory.load_memory_variables({})['history'], combined_context)
+
+        self.update_chat(chat_id, question, answer)
 
-        
-        
         if language == 'urdu':
-            translated_answer = self.translate_urdu(answer)
-            self.update_chat(chat_id, question, answer)
-            return translated_answer
-        else:
-            self.update_chat(chat_id, question, answer)
-            return answer
+            return self.translate_urdu(answer)
+        return answer
+                      
+