Spaces:
Sleeping
Sleeping
| # 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 | |
| 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])}" # 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(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}) # 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 | |
| ) |