Update application/chat_inference.py

application/chat_inference.py

@@ -1,70 +1,93 @@
+# application/chat_inference.py
+import time  # Import the time module
 from application.utils.chat_completion_api import ChatCompletionAPI
-from config import Response,pipeline_dict,convs_dict
+from config import Response, pipeline_dict, convs_dict
 import os
 from application.utils.image_captioning import ImageCaptioning
+from application.utils.web_search import WebScarper
+from application.utils.image_generation import generate_image  # Import
+
 
 class ChatInference:
     def __init__(self):
         self.chatCompletionAPI = ChatCompletionAPI()
         self.image_captioning = ImageCaptioning()
+        self.web_scraper = WebScarper()
 
-    def validate(self,data,user):
-        try:
-            pipeline = pipeline_dict['api']['models']
-            model = data['model']
-            self.headers = pipeline[model]['headers']
-            self.updateHeaders = {}
-            for header in self.headers:
-                if(header=="config"):
-                    for configHeader in self.headers[header]:
-                        if(configHeader=="Authorization"):
-                            auth = self.headers[header][configHeader].split(' ')
-                            self.updateHeaders[configHeader] = f"{auth[0]} {eval(auth[1])}"
-                        elif(configHeader=="comment"):
-                            pass
-                        else:
-                            self.updateHeaders[configHeader] = f"{eval(self.headers[header][configHeader])}"
-                else:
-                    self.updateHeaders[header] = self.headers[header]
-            prompt = data['prompt']
-            max_tokens = data.get('max_token', 10020)
-            temperature = max(0, min(data.get('temperature', 0.7), 2))
-            top_p = max(0.1, min(data.get('top_p', 0.9), 1))
-            system = data.get('system_prompt','You are a helpful and harmless AI assistant. You are xylaria made by sk md saad amin. You should think step-by-step')
-            convId = data['convId']
-            image = data.get('image')
 
-            if(len(convs_dict[user][convId]['messages'])==1):
-                if system:
-                    convs_dict[user][convId]['messages'].insert(0,{"role":"system", "content": system})
-                convs_dict[user]['metadata'].insert(0,{"convId": convId, "title": prompt[:23]})
-                convs_dict[user][convId]['title'] = prompt[:30]
+    def validate(self, data, user):
+      try:
+        pipeline = pipeline_dict['api']['models']
+        model = data['model']
+        self.headers = pipeline[model]['headers']
+        self.updateHeaders = {}
+        for header in self.headers:
+            if(header=="config"):
+                for configHeader in self.headers[header]:
+                    if(configHeader=="Authorization"):
+                        auth = self.headers[header][configHeader].split(' ')
+                        self.updateHeaders[configHeader] = f"{auth[0]} {eval(auth[1])}"  # Directly evaluate
+                    elif(configHeader=="comment"):
+                        pass
+                    else:
+                        self.updateHeaders[configHeader] = f"{eval(self.headers[header][configHeader])}" # Directly evaluate
+            else:
+                self.updateHeaders[header] = self.headers[header]
+        prompt = data['prompt']
+        max_tokens = data.get('max_token', 10020)
+        temperature = max(0, min(data.get('temperature', 0.7), 2))
+        top_p = max(0.1, min(data.get('top_p', 0.9), 1))
+        system = data.get('system_prompt','You are a helpful and harmless AI assistant. You are xylaria made by sk md saad amin. You should think step-by-step')
+        convId = data['convId']
+        image = data.get('image')
 
-            if image:
-                caption = self.image_captioning.generate_caption(image)
-                prompt = f"{caption}\n\n{prompt}"
+        if(len(convs_dict[user][convId]['messages'])==1):
+            if system:
+                # Include user memory in the system prompt
+                system_prompt = f"{system}\n\nMemory: {convs_dict[user]['memory']}"
+                convs_dict[user][convId]['messages'][0]['content'] = system_prompt  # Update existing system message
 
+            convs_dict[user]['metadata'].insert(0,{"convId": convId, "title": prompt[:23]})
+            convs_dict[user][convId]['title'] = prompt[:30]
+        if image:
+            caption = self.image_captioning.generate_caption(image)
+            prompt = f"{caption}\n\n{prompt}"
 
-            if(pipeline[model]['type'] == 'image-text-to-text'):
-                convs_dict[user][convId]['messages'].append({"role": "user", "content": [{"type":"text","text":prompt}]})
-            else:
-                convs_dict[user][convId]['messages'].append({"role":"user","content":prompt})
-            transformed = {
-                "model": model,
-                "prompt": prompt,
-                "messages": convs_dict[user][convId]['messages'],
-                "max_tokens": max_tokens,
-                "temperature": temperature,
-                "top_p": top_p,
-                "stream": True
-            }
-            data.update(transformed)
-            return data
-        except KeyError:
-            return 400
 
-    def chat(self,data,handle_stream,user):
-        data = self.validate(data=data,user=user)
-        if(data==400):
-            return "Required Parameters are Missing!", 400
-        return self.chatCompletionAPI.make_request(json=data,url=data['base_url'],handle_stream=handle_stream,messages=data['messages'], headers=self.updateHeaders, webSearch=data['webSearch'])
+        if(pipeline[model]['type'] == 'image-text-to-text'):
+            convs_dict[user][convId]['messages'].append({"role": "user", "content": [{"type":"text","text":prompt}]})
+        else:
+            convs_dict[user][convId]['messages'].append({"role":"user","content":prompt})  # Append user message
+        transformed = {
+            "model": model,
+            "prompt": prompt,
+            "messages": convs_dict[user][convId]['messages'],
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "stream": True
+        }
+        data.update(transformed)
+        return data
+      except KeyError as e:
+        print(f"KeyError: {e}")  # Debugging
+        return 400
+      except Exception as e:
+        print(f"An unexpected error occurred: {e}")  # Debugging
+        return 500
+
+    def chat(self, data, handle_stream, user):
+        start_time = time.time()  # Capture start time
+        data = self.validate(data=data, user=user)
+        if isinstance(data, int):  # Check for error codes
+            return "Required Parameters are Missing!", data
+
+        return self.chatCompletionAPI.make_request(
+            json=data,
+            url=data['base_url'],
+            handle_stream=handle_stream,
+            messages=data['messages'],
+            headers=self.updateHeaders,
+            webSearch=data['webSearch'],
+            start_time = start_time
+        )