Spaces:
Running
Running
| from backend import get_handler | |
| import uuid | |
| import os | |
| import json | |
| import time | |
| from datetime import datetime | |
| from app_utils import * | |
| class State: | |
| """ | |
| Manages the state of a chatbot including its model configuration, chat history, and test data. | |
| Attributes: | |
| current_model (str): The currently active model identifier | |
| current_temperature (float): The temperature setting for model inference | |
| current_category (str): The current category of conversation | |
| test_entry (dict): Contains test configuration and data | |
| inference_data (dict): Stores inference results and messages | |
| handler: The model handler instance | |
| history (list): List of conversation messages | |
| name (str): Identifier for this bot instance, defaults to "bot" | |
| database: Database connection for storing chat history | |
| example_settings (dict): Predefined example configurations (see app_utils.py) | |
| """ | |
| def __init__(self, init_model, init_temperature, database, example_settings, name="bot"): | |
| self.current_model = init_model | |
| self.current_temperature = init_temperature | |
| self.current_category = None | |
| self.test_entry = initialize_empty_test_entry() | |
| self.inference_data = {"message": []} | |
| self.handler = get_handler(self.current_model, self.current_temperature) | |
| self.history = self.get_initial_state() | |
| self.name = name | |
| self.database = database | |
| self.example_settings = example_settings | |
| def initialize(self, category): | |
| self.category = category | |
| self.update_category_and_load_config(category) | |
| def get_initial_state(self): | |
| return list(INITIAL_CHAT_HISTORY) | |
| def restart_chat(self, model, temperature): | |
| self.test_entry["id"] = str(uuid.uuid4()) | |
| self.handler = get_handler(model, temperature) | |
| self.history = self.get_initial_state() | |
| return self.history | |
| def update_handler(self, model, temperature): | |
| self.current_model = model | |
| self.current_temperature = temperature | |
| self.handler = get_handler(model, temperature) | |
| print(f"Update handler for {self.name}: ", model, temperature) | |
| self.history = self.restart_chat(model, temperature) | |
| return model, self.history | |
| def save(self): | |
| if len(self.history) > 2: | |
| document = {"time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), | |
| "model": self.current_model, | |
| "category": self.current_category, | |
| "temperature": self.current_temperature, | |
| "history": self.history} | |
| self.database.insert_one(document) | |
| def restart_chat_and_save(self): | |
| self.save() | |
| return self.restart_chat(self.current_model, self.current_temperature) | |
| def update_category_and_load_config(self, category): | |
| self.current_category = category | |
| self.test_entry["initial_config"] = {category: {}} | |
| self.test_entry["involved_classes"] = [category] | |
| config_path = os.path.join("config", f"{MAPPINGS[category]}.json") | |
| self.load_config(config_path) | |
| return category | |
| def load_config(self, config_path): | |
| if os.path.exists(config_path): | |
| with open(config_path, 'r') as config_file: | |
| data = json.load(config_file) | |
| self.test_entry["function"] = data.copy() | |
| def load_example_and_update(self, example): | |
| self.save() | |
| model, temp, category, message = self.load_example(example) | |
| self.update_category_and_load_config(category) | |
| return model, temp, category, message | |
| def load_example(self, example): | |
| return self.example_settings[example] | |
| def response(self): | |
| for item in self.handler.inference(self.test_entry): | |
| if item[0] == "regular": | |
| responses_results = equalize_and_zip(item[1], item[2]) | |
| for (model_res, exec_res) in responses_results: | |
| if model_res is not None: | |
| response = model_res | |
| self.history.append({"role": "assistant", "content": "<b>Model Response🤖: </b><br>"}) | |
| for character in response: | |
| self.history[-1]["content"] += character | |
| time.sleep(0.01) | |
| yield self.history | |
| if exec_res is not None: | |
| response = exec_res | |
| self.history[-1]["content"] += "<br><br><b>Model Execution💻: </b><br>" | |
| yield self.history | |
| for character in response: | |
| self.history[-1]["content"] += character | |
| time.sleep(0.01) | |
| yield self.history | |
| elif item[0] == 'summary': | |
| response = item[1] | |
| if response is not None: | |
| self.history.append({"role": "assistant", "content": "<b>Summary✅: </b><br>"}) | |
| for character in response: | |
| self.history[-1]["content"] += character | |
| time.sleep(0.01) | |
| yield self.history | |
| elif item[0] == "final": | |
| self.inference_data = item[2] | |
| time.sleep(0.05) |