Update app.py

app.py

@@ -7,9 +7,8 @@ from datetime import datetime
 import logging
 
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from huggingface_hub import InferenceClient, cached_download
-from safe_search import safe_search
-from i_search import google, i_search as i_s
 
 # --- Configuration ---
 VERBOSE = True  # Enable verbose logging
@@ -23,8 +22,8 @@ API_KEY = "YOUR_API_KEY"  # Replace with your actual Hugging Face API key
 
 # --- Logging Setup ---
 logging.basicConfig(
-    filename="app.log",  # Name of the log file
-    level=logging.INFO,  # Set the logging level (INFO, DEBUG, etc.)
+    filename="app.log",
+    level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s",
 )
 
@@ -41,7 +40,7 @@ agents = [
 PREFIX = """
 {date_time_str}
 Purpose: {purpose}
-Safe Search: {safe_search}
+Agent: {agent_name}
 """
 
 LOG_PROMPT = """
@@ -52,70 +51,14 @@ 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: 
-"""
-
 # --- Functions ---
 def format_prompt(message: str, history: List[Tuple[str, str]], max_history_turns: int = 2) -> str:
-    """Formats the prompt for the LLM, including the message and relevant history."""
-    prompt = " "
-    # Keep only the last 'max_history_turns' turns
+    prompt = ""
     for user_prompt, bot_response in history[-max_history_turns:]:
-        prompt += f"[INST] {user_prompt} [/  "
-        prompt += f" {bot_response}"
-    prompt += f"[INST] {message} [/  "
+        prompt += f"Human: {user_prompt}\nAssistant: {bot_response}\n"
+    prompt += f"Human: {message}\nAssistant:"
     return prompt
 
-def run_llm(
-    prompt_template: str,
-    stop_tokens: List[str],
-    purpose: str,
-    **prompt_kwargs: Dict
-) -> str:
-    """Runs the LLM with the given prompt and parameters."""
-    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=TEMPERATURE, top_p=TOP_P, repetition_penalty=REPETITION_PENALTY)
-    if VERBOSE:
-        logging.info(LOG_RESPONSE.format(resp))  # Log the response
-    return resp
-
 def generate(
     prompt: str,
     history: List[Tuple[str, str]],
@@ -126,206 +69,43 @@ def generate(
     top_p: float = TOP_P,
     repetition_penalty: float = REPETITION_PENALTY,
 ) -> str:
-    """Generates text using the LLM."""
-    content = PREFIX.format(
+    # Load model and tokenizer
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    
+    # Create a text generation pipeline
+    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    
+    # Prepare the full prompt
+    date_time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    full_prompt = 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
-
-    stream = client.text_generation(content, stream=True, details=True, return_full_text=False, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty, max_new_tokens=max_new_tokens)
-    resp = ""
-    for response in stream:
-        resp += response.token.text
-
+        purpose=sys_prompt,
+        agent_name=agent_name
+    ) + format_prompt(prompt, history)
+    
     if VERBOSE:
-        logging.info(LOG_RESPONSE.format(resp))  # Log the response
-    return resp
-
-def compress_history(purpose: str, task: str, history: List[Tuple[str, str]], directory: str) -> str:
-    """Compresses the history into a shorter summary."""
-    resp = run_llm(
-        COMPRESS_HISTORY_PROMPT,
-        stop_tokens=["observation:", "task:", "action:", "thought:"],
-        purpose=purpose,
-        task=task,
-        history="\n".join(f"[INST] {user_prompt} [/] {bot_response}" for user_prompt, bot_response in history),
-    )
-    history = "observation: {}\n".format(resp)
-    return history
-
-def call_search(purpose: str, task: str, history: List[Tuple[str, str]], directory: str, action_input: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
-    """Performs a search based on the 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.append(("observation: search result is: {}".format(response), ""))
-        else:
-            history.append(("observation: I need to provide a valid URL to 'action: SEARCH action_input=https://URL'\n", ""))
-    except Exception as e:
-        history.append(("observation: {}\n".format(e), ""))
-    return "MAIN", None, history, task
-
-def call_main(purpose: str, task: str, history: List[Tuple[str, str]], directory: str, action_input: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
-    """Handles the main agent interaction loop."""
-    logging.info(f"CALLING MAIN: {action_input}")
-    resp = run_llm(
-        ACTION_PROMPT,
-        stop_tokens=["observation:", "task:", "action:", "thought:"],
-        purpose=purpose,
-        task=task,
-        history="\n".join(f"[INST] {user_prompt} [/] {bot_response}" for user_prompt, bot_response in history),
-    )
-    lines = resp.strip().strip("\n").split("\n")
-    for line in lines:
-        if line == "":
-            continue
-        if line.startswith("thought: "):
-            history.append((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.append((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.append((line, ""))
-            logging.info(f"Other Output: {line}")
-    return "MAIN", None, history, task
-
-def call_set_task(purpose: str, task: str, history: List[Tuple[str, str]], directory: str, action_input: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
-    """Sets a new task for the agent."""
-    logging.info(f"CALLING SET_TASK: {action_input}")
-    task = run_llm(
-        TASK_PROMPT,
-        stop_tokens=[],
-        purpose=purpose,
-        task=task,
-        history="\n".join(f"[INST] {user_prompt} [/] {bot_response}" for user_prompt, bot_response in history),
-    ).strip("\n")
-    history.append(("observation: task has been updated to: {}".format(task), ""))
-    return "MAIN", None, history, task
-
-def end_fn(purpose: str, task: str, history: List[Tuple[str, str]], directory: str, action_input: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
-    """Ends the agent interaction."""
-    logging.info(f"CALLING END_FN: {action_input}")
-    task = "END"
-    return "COMPLETE", "COMPLETE", history, task
-
-NAME_TO_FUNC: Dict[str, callable] = {
-    "MAIN": call_main,
-    "UPDATE-TASK": call_set_task,
-    "SEARCH": call_search,
-    "COMPLETE": end_fn,
-}
-
-def run_action(purpose: str, task: str, history: List[Tuple[str, str]], directory: str, action_name: str, action_input: str) -> Tuple[str, str, List[Tuple[str, str]], str]:
-    """Executes the specified action."""
-    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
+        logging.info(LOG_PROMPT.format(content=full_prompt))
     
-        # compress the history when it is long
-        if len(history) > 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 is None:
-            action_name = "MAIN"
-        assert action_name in NAME_TO_FUNC
+    # Generate response
+    response = generator(
+        full_prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True
+    )[0]['generated_text']
     
-        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.append(("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: str, history: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
-    """Main agent interaction loop."""
-    task = None
-    directory = "./"
-    if history:
-        history = str(history).strip("[]")
-    if not history:
-        history = []
+    # Extract the assistant's response
+    assistant_response = response.split("Assistant:")[-1].strip()
     
-    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 format_prompt(message: str, history: List[Tuple[str, str]], max_history_turns: int = 5) -> str:
-    """Formats the prompt for the LLM, including the message and relevant history."""
-    prompt = " "
-    # Keep only the last 'max_history_turns' turns
-    for user_prompt, bot_response in history[-max_history_turns:]:
-        prompt += f"[INST] {user_prompt} [/  "
-        prompt += f" {bot_response}"
-    prompt += f"[INST] {message} [/  "
-    return prompt
-
-def parse_action(line: str) -> Tuple[str, str]:
-    """Parses the action line to get the action name and input."""
-    parts = line.split(":", 1)
-    if len(parts) == 2:
-        action_name = parts[0].replace("action", "").strip()
-        action_input = parts[1].strip()
-    else:
-        action_name = parts[0].replace("action", "").strip()
-        action_input = ""
-    return action_name, action_input
+    if VERBOSE:
+        logging.info(LOG_RESPONSE.format(resp=assistant_response))
+    
+    return assistant_response
 
 def main():
-    """Main function to run the Gradio interface."""
-    global client
-    # Initialize the LLM client with your API key
-    try:
-        client = InferenceClient(
-            MODEL_NAME,
-            token=API_KEY  # Replace with your actual API key
-        )
-    except Exception as e:
-        logging.error(f"Error initializing LLM client: {e}")
-        print("Error initializing LLM client. Please check your API key.")
-        return
-
     with gr.Blocks() as demo:
         gr.Markdown("## FragMixt: The No-Code Development Powerhouse")
         gr.Markdown("###  Your AI-Powered Development Companion")
@@ -365,8 +145,7 @@ def main():
 
         def chat(purpose: str, message: str, agent_name: str, sys_prompt: str, temperature: float, max_new_tokens: int, top_p: float, repetition_penalty: float, history: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]]]:
             """Handles the chat interaction."""
-            prompt = format_prompt(message, history)
-            response = generate(prompt, history, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            response = generate(message, history, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
             history.append((message, response))
             return history, history
 
@@ -386,4 +165,5 @@ def main():
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()
+```