Spaces:
Sleeping
Sleeping
| import os | |
| import subprocess | |
| import random | |
| from huggingface_hub import InferenceClient | |
| import gradio as gr | |
| from safe_search import safe_search | |
| from i_search import google | |
| from i_search import i_search as i_s | |
| from datetime import datetime | |
| import logging | |
| import json | |
| import nltk | |
| from transformers import pipeline | |
| # Ensure NLTK data is downloaded | |
| nltk.download('punkt') | |
| now = datetime.now() | |
| date_time_str = now.strftime("%Y-%m-%d %H:%M:%S") | |
| client = InferenceClient( | |
| "mistralai/Mixtral-8x7B-Instruct-v0.1" | |
| ) | |
| # --- Set up logging --- | |
| logging.basicConfig( | |
| filename="app.log", # Name of the log file | |
| level=logging.INFO, # Set the logging level (INFO, DEBUG, etc.) | |
| format="%(asctime)s - %(levelname)s - %(message)s", | |
| ) | |
| agents = [ | |
| "WEB_DEV", | |
| "AI_SYSTEM_PROMPT", | |
| "PYTHON_CODE_DEV" | |
| ] | |
| VERBOSE = True | |
| MAX_HISTORY = 5 | |
| PREFIX = """ | |
| {date_time_str} | |
| Purpose: {purpose} | |
| Safe Search: {safe_search} | |
| """ | |
| LOG_PROMPT = """ | |
| PROMPT: {content} | |
| """ | |
| LOG_RESPONSE = """ | |
| RESPONSE: {resp} | |
| """ | |
| COMPRESS_HISTORY_PROMPT = """ | |
| You are a helpful AI assistant. Your task is to compress the following history into a summary that is no longer than 512 tokens. | |
| History: | |
| {history} | |
| """ | |
| ACTION_PROMPT = """ | |
| You are a helpful AI assistant. You are working on the task: {task} | |
| Your current history is: | |
| {history} | |
| What is your next thought? | |
| thought: | |
| What is your next action? | |
| action: | |
| """ | |
| TASK_PROMPT = """ | |
| You are a helpful AI assistant. Your current history is: | |
| {history} | |
| What is the next task? | |
| task: | |
| """ | |
| UNDERSTAND_TEST_RESULTS_PROMPT = """ | |
| You are a helpful AI assistant. The test results are: | |
| {test_results} | |
| What do you want to know about the test results? | |
| thought: | |
| """ | |
| def format_prompt(message, history, max_history_turns=5): | |
| prompt = "<s>" | |
| # Keep only the last 'max_history_turns' turns | |
| for user_prompt, bot_response in history[-max_history_turns:]: | |
| prompt += f"[INST] {user_prompt} [/INST]" | |
| prompt += f" {bot_response}</s> " | |
| prompt += f"[INST] {message} [/INST]" | |
| return prompt | |
| def run_gpt( | |
| prompt_template, | |
| stop_tokens, | |
| max_tokens, | |
| purpose, | |
| **prompt_kwargs, | |
| ): | |
| seed = random.randint(1,1111111111111111) | |
| logging.info(f"Seed: {seed}") # Log the seed | |
| content = PREFIX.format( | |
| date_time_str=date_time_str, | |
| purpose=purpose, | |
| safe_search=safe_search, | |
| ) + prompt_template.format(**prompt_kwargs) | |
| if VERBOSE: | |
| logging.info(LOG_PROMPT.format(content)) # Log the prompt | |
| resp = client.text_generation(content, max_new_tokens=max_tokens, stop_sequences=stop_tokens, temperature=0.7, top_p=0.8, repetition_penalty=1.5) | |
| if VERBOSE: | |
| logging.info(LOG_RESPONSE.format(resp=resp)) # Log the response | |
| return resp | |
| def generate( | |
| prompt, history, agent_name=agents[0], sys_prompt="", temperature=0.7, max_new_tokens=2048, top_p=0.8, repetition_penalty=1.5, model="mistralai/Mixtral-8x7B-Instruct-v0.1" | |
| ): | |
| seed = random.randint(1,1111111111111111) | |
| if agent_name == "WEB_DEV": | |
| agent = "You are a helpful AI assistant. You are a web developer." | |
| elif agent_name == "AI_SYSTEM_PROMPT": | |
| agent = "You are a helpful AI assistant. You are an AI system." | |
| elif agent_name == "PYTHON_CODE_DEV": | |
| agent = "You are a helpful AI assistant. You are a Python code developer." | |
| else: | |
| agent = "You are a helpful AI assistant." | |
| system_prompt = agent | |
| temperature = float(temperature) | |
| if temperature < 1e-2: | |
| temperature = 1e-2 | |
| top_p = float(top_p) | |
| formatted_prompt = f"{system_prompt} {prompt}" | |
| formatted_prompt = format_prompt(formatted_prompt, history, max_history_turns=5) # Truncated history | |
| logging.info(f"Formatted Prompt: {formatted_prompt}") | |
| stream = client.text_generation(formatted_prompt, temperature=temperature, max_new_tokens=max_new_tokens, top_p=top_p, repetition_penalty=repetition_penalty, stream=True, details=True, return_full_text=False) | |
| resp = "" | |
| for response in stream: | |
| resp += response.token.text | |
| if VERBOSE: | |
| logging.info(LOG_RESPONSE.format(resp=resp)) # Log the response | |
| return resp | |
| def compress_history(purpose, task, history, directory): | |
| resp = run_gpt( | |
| COMPRESS_HISTORY_PROMPT, | |
| stop_tokens=["observation:", "task:", "action:", "thought:"], | |
| max_tokens=512, | |
| purpose=purpose, | |
| task=task, | |
| history=history, | |
| ) | |
| history = "observation: {}\n".format(resp) | |
| return history | |
| def call_search(purpose, task, history, directory, action_input): | |
| logging.info(f"CALLING SEARCH: {action_input}") | |
| try: | |
| if "http" in action_input: | |
| if "<" in action_input: | |
| action_input = action_input.strip("<") | |
| if ">" in action_input: | |
| action_input = action_input.strip(">") | |
| response = i_s(action_input) | |
| logging.info(f"Search Result: {response}") | |
| history += "observation: search result is: {}\n".format(response) | |
| else: | |
| history += "observation: I need to provide a valid URL to 'action: SEARCH action_input=https://URL'\n" | |
| except Exception as e: | |
| history += "observation: {}'\n".format(e) | |
| return "MAIN", None, history, task | |
| def call_main(purpose, task, history, directory, action_input): | |
| logging.info(f"CALLING MAIN: {action_input}") | |
| resp = run_gpt( | |
| ACTION_PROMPT, | |
| stop_tokens=["observation:", "task:", "action:","thought:"], | |
| max_tokens=32000, | |
| purpose=purpose, | |
| task=task, | |
| history=history, | |
| ) | |
| lines = resp.strip().strip("\n").split("\n") | |
| for line in lines: | |
| if line == "": | |
| continue | |
| if line.startswith("thought: "): | |
| history += "{}\n".format(line) | |
| logging.info(f"Thought: {line}") | |
| elif line.startswith("action: "): | |
| action_name, action_input = parse_action(line) | |
| logging.info(f"Action: {action_name} - {action_input}") | |
| history += "{}\n".format(line) | |
| if "COMPLETE" in action_name or "COMPLETE" in action_input: | |
| task = "END" | |
| return action_name, action_input, history, task | |
| else: | |
| return action_name, action_input, history, task | |
| else: | |
| history += "{}\n".format(line) | |
| logging.info(f"Other Output: {line}") | |
| return "MAIN", None, history, task | |
| def call_set_task(purpose, task, history, directory, action_input): | |
| logging.info(f"CALLING SET_TASK: {action_input}") | |
| task = run_gpt( | |
| TASK_PROMPT, | |
| stop_tokens=[], | |
| max_tokens=64, | |
| purpose=purpose, | |
| task=task, | |
| history=history, | |
| ).strip("\n") | |
| history += "observation: task has been updated to: {}\n".format(task) | |
| return "MAIN", None, history, task | |
| def end_fn(purpose, task, history, directory, action_input): | |
| logging.info(f"CALLING END_FN: {action_input}") | |
| task = "END" | |
| return "COMPLETE", "COMPLETE", history, task | |
| NAME_TO_FUNC = { | |
| "MAIN": call_main, | |
| "UPDATE-TASK": call_set_task, | |
| "SEARCH": call_search, | |
| "COMPLETE": end_fn, | |
| } | |
| def run_action(purpose, task, history, directory, action_name, action_input): | |
| logging.info(f"RUNNING ACTION: {action_name} - {action_input}") | |
| try: | |
| if "RESPONSE" in action_name or "COMPLETE" in action_name: | |
| action_name="COMPLETE" | |
| task="END" | |
| return action_name, "COMPLETE", history, task | |
| # compress the history when it is long | |
| if len(history.split("\n")) > MAX_HISTORY: | |
| logging.info("COMPRESSING HISTORY") | |
| history = compress_history(purpose, task, history, directory) | |
| if not action_name in NAME_TO_FUNC: | |
| action_name="MAIN" | |
| if action_name == "" or action_name == None: | |
| action_name="MAIN" | |
| assert action_name in NAME_TO_FUNC | |
| logging.info(f"RUN: {action_name} - {action_input}") | |
| return NAME_TO_FUNC[action_name](purpose, task, history, directory, action_input) | |
| except Exception as e: | |
| history += "observation: the previous command did not produce any useful output, I need to check the commands syntax, or use a different command\n" | |
| logging.error(f"Error in run_action: {e}") | |
| return "MAIN", None, history, task | |
| def run(purpose, history): | |
| task = None | |
| directory = "./" | |
| if history: | |
| history = str(history).strip("[]") | |
| if not history: | |
| history = "" | |
| action_name = "UPDATE-TASK" if task is None else "MAIN" | |
| action_input = None | |
| while True: | |
| logging.info(f"---") | |
| logging.info(f"Purpose: {purpose}") | |
| logging.info(f"Task: {task}") | |
| logging.info(f"---") | |
| logging.info(f"History: {history}") | |
| logging.info(f"---") | |
| action_name, action_input, history, task = run_action( | |
| purpose, | |
| task, | |
| history, | |
| directory, | |
| action_name, | |
| action_input, | |
| ) | |
| yield (history) | |
| if task == "END": | |
| return (history) | |
| def generate_text_chunked(input_text, model, generation_parameters, max_tokens_to_generate): | |
| """Generates text in chunks to avoid token limit errors.""" | |
| sentences = nltk.sent_tokenize(input_text) | |
| generated_text = [] | |
| generator = pipeline('text-generation', model=model) | |
| for sentence in sentences: | |
| # Tokenize the sentence and check if it's within the limit | |
| tokens = generator.tokenizer(sentence).input_ids | |
| if len(tokens) + max_tokens_to_generate <= 32768: | |
| # Generate text for this chunk | |
| response = generator(sentence, max_length=max_tokens_to_generate, **generation_parameters) | |
| generated_text.append(response[0]['generated_text']) | |
| else: | |
| # Handle cases where the sentence is too long | |
| print(f"Sentence too long: {sentence}") | |
| return ''.join(generated_text) | |
| # Gradio Interface | |
| def gradio_interface(purpose, history): | |
| try: | |
| history = json.loads(history) if history else [] | |
| except json.JSONDecodeError: | |
| history = [] | |
| result = run(purpose, history) | |
| return next(result) | |
| iface = gr.Interface( | |
| fn=gradio_interface, | |
| inputs=[ | |
| gr.Textbox(lines=2, placeholder="Enter the purpose here..."), | |
| gr.Textbox(lines=10, placeholder="Enter the history here (JSON format)...") | |
| ], | |
| outputs="text", | |
| title="AI Assistant", | |
| description="An AI assistant that helps with various tasks." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() |