Spaces:
Build error
Build error
| import os,re,sys,yaml | |
| from rich import print as rp | |
| from dotenv import load_dotenv, find_dotenv | |
| from improvement_prompts import first_prompt, second_prompt,new_first_prompt, new_second_prompt | |
| load_dotenv(find_dotenv()) | |
| import logging | |
| from hugchat import hugchat | |
| from hugchat.login import Login | |
| from typing import List, Dict, Generator, Optional,Tuple | |
| logging.basicConfig(filename='improvement_log.log', level=logging.INFO) | |
| from UberToolkit import UberToolkit as UberK | |
| class HugChatLLM: | |
| def __init__(self, cookie_path_dir: str = "./cookies/"): | |
| self.email = os.getenv("EMAIL") | |
| self.password = os.getenv("PASSWD") | |
| if not self.email or not self.password: | |
| print("EMAIL and PASSWD environment variables must be set.") | |
| sys.exit(1) | |
| self.cookie_path_dir = cookie_path_dir | |
| self.chatbot = self._login_and_create_chatbot() | |
| self.current_conversation = self.chatbot.new_conversation(modelIndex=1, system_prompt='', switch_to=True) | |
| self.img_gen_servers = ["Yntec/HuggingfaceDiffusion", "Yntec/WinningBlunder"] | |
| #self.igllm=ImageGeneratorLLM() | |
| self.triggers={'/improve:' :'Improve the following text: ', | |
| '/fix:' :'Fix the following code: ', | |
| "/save_yaml:" :self.save_yaml, | |
| "/add_tab:" :"", | |
| "/edit_tab:" :"", | |
| } | |
| def _login_and_create_chatbot(self) -> hugchat.ChatBot: | |
| sign = Login(self.email, self.password) | |
| cookies = sign.login(cookie_dir_path=self.cookie_path_dir, save_cookies=True) | |
| return hugchat.ChatBot(cookies=cookies.get_dict()) | |
| def __call__(self, text: str, stream: bool = True, web_search: bool = False): | |
| if stream: | |
| return self.trigger_mon(input_string=self.query(text)) | |
| else: | |
| text_result, code_result = self.trigger_mon(self.query(text, web_search)) | |
| return text_result, code_result | |
| def query(self, text: str, web_search: bool = False) -> str: | |
| if not text == '': | |
| message_result = self.chatbot.chat(text, web_search=web_search) | |
| return message_result.wait_until_done() | |
| def stream_query(self, text: str) -> Generator[str, None, None]: | |
| if not text == '': | |
| for resp in self.chatbot.chat(text, stream=True): | |
| yield resp | |
| def new_conversation(self, model_index=1, system_prompt='', switch_to: bool = True) -> str: | |
| return self.chatbot.new_conversation(modelIndex=model_index, system_prompt=system_prompt, switch_to=switch_to) | |
| #return self.chatbot.get_conversation_info() | |
| def get_remote_conversations(self, replace_conversation_list: bool = True) -> List[Dict]: | |
| return self.chatbot.get_remote_conversations(replace_conversation_list=replace_conversation_list) | |
| def get_conversation_list(self) -> List[Dict]: | |
| return self.chatbot.get_conversation_list() | |
| def get_available_models(self) -> List[str]: | |
| return self.chatbot.get_available_llm_models() | |
| def switch_model(self, index: int) -> None: | |
| self.chatbot.switch_llm(index) | |
| def get_conversation_info(self) -> Dict: | |
| info = self.chatbot.get_conversation_info() | |
| return { | |
| "id": info.id, | |
| "title": info.title, | |
| "model": info.model, | |
| "system_prompt": info.system_prompt, | |
| "history": info.history | |
| } | |
| def search_assistant_by_name(self, assistant_name: str) -> Dict: | |
| return self.chatbot.search_assistant(assistant_name=assistant_name) | |
| def search_assistant_by_id(self, assistant_id: str) -> Dict: | |
| return self.chatbot.search_assistant(assistant_id=assistant_id) | |
| def get_assistant_list(self, page: int = 0) -> List[Dict]: | |
| return self.chatbot.get_assistant_list_by_page(page=page) | |
| def new_conversation_with_assistant(self, assistant: Dict, switch_to: bool = True) -> None: | |
| self.chatbot.new_conversation(assistant=assistant, switch_to=switch_to) | |
| return self.chatbot.get_conversation_info() | |
| def delete_all_conversations(self) -> None: | |
| self.chatbot.delete_all_conversations() | |
| def get_available_models(self) -> List[Dict[str, str]]: | |
| models = self.chatbot.get_available_llm_models() | |
| for model in models: | |
| logging.info(model.id) | |
| logging.info(model.name) | |
| return [{"id": model.id,"name": model.name} for model in models] | |
| def save_yaml(self, input_string: str, file_path: str, file_name: str) -> str: | |
| yaml_file_path = os.path.join(file_path, file_name) | |
| rp.print(f'Saving YAML data: {input_string}\n\nTo: {yaml_file_path}') | |
| # here we write the YAML data to a file | |
| with open(yaml_file_path, 'w') as file: | |
| yaml.dump(input_string, file) | |
| return "YAML data saved successfully." | |
| def trigger_mon(self, input_string) -> str: | |
| '''in this we need to detect if the | |
| AI response contains a key of the Dict triggers ers then fetch the input directly from the response behind the the triggerger.... | |
| and call the method lidyed in the value in the triggers Dict''' | |
| # here we detect and route the trigger keys from responses | |
| # using a dictionary mapping trigger keys to their respective prompts | |
| for trigger, action in self.triggers.items(): | |
| if trigger in input_string: | |
| input_split = input_string.split(trigger).pop().split('\n').pop(0) | |
| rest_split = ' '.join(input_string.split(f'{trigger}{input_split}').replace('input','[trigger_processed]input')) | |
| print(f'Detected trigger: {trigger}, fetching input: {input_split}') | |
| action(input_split) | |
| # if the trigger is found, we need to fetch the input directly from the response | |
| return rest_split | |
| # if no trigger is found, we return the input string as is | |
| return input_string | |
| class ImageGeneratorLLM(HugChatLLM): | |
| def __init__(self): | |
| super().__init__() | |
| def __call__(self, prompt: str) -> str: | |
| return self.generate_image(prompt) | |
| def generate_image(self, prompt: str) -> str: | |
| return self.chatbot.generate_image(prompt) |