app.py

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

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
#MODEL = "gpt-3.5-turbo"  # "gpt-4"

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=2):
  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))  # 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):
    seed = random.randint(1,1111111111111111)

    # Correct the line:
    if agent_name == "WEB_DEV":
        agent = "You are a helpful AI assistant. You are a web developer."
    if agent_name == "AI_SYSTEM_PROMPT":
        agent = "You are a helpful AI assistant. You are an AI system."
    if agent_name == "PYTHON_CODE_DEV":
        agent = "You are a helpful AI assistant. You are a Python code developer."
    system_prompt = agent
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=seed,
    )

    formatted_prompt = format_prompt(prompt, history, max_history_turns=5)  # Truncated history
    logging.info(f"Formatted Prompt: {formatted_prompt}")

    messages = [{"role": "user", "content": formatted_prompt}]

    stream = client.text_generation(messages, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    for response in stream:
        output += response.token.text
        yield output
    return output


additional_inputs=[
    gr.Dropdown(
        label="Agents",
        choices=[s for s in agents],
        value=agents[0],
        interactive=True,
        ),
    gr.Textbox(
        label="System Prompt",
        max_lines=1,
        interactive=True,
    ),
    gr.Slider(
        label="Temperature",
        value=0.9,
        minimum=0.0,
        maximum=1.0,
        step=0.05,
        interactive=True,
        info="Higher values produce more diverse outputs",
    ),

    gr.Slider(
        label="Max new tokens",
        value=1048*10,
        minimum=0,
        maximum=1048*10,
        step=64,
        interactive=True,
        info="The maximum numbers of new tokens",
    ),
    gr.Slider(
        label="Top-p (nucleus sampling)",
        value=0.90,
        minimum=0.0,
        maximum=1,
        step=0.05,
        interactive=True,
        info="Higher values sample more low-probability tokens",
    ),
    gr.Slider(
        label="Repetition penalty",
        value=1.2,
        minimum=1.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Penalize repeated tokens",
    ),


]

examples = [
    ["Help me set up TypeScript configurations and integrate ts-loader in my existing React project.",
"Update Webpack Configurations",
"Install Dependencies",
"Configure Ts-Loader",
"TypeChecking Rules Setup",
"React Specific Settings",
"Compilation Options",
"Test Runner Configuration"],

["Guide me through building a serverless microservice using AWS Lambda and API Gateway, connecting to DynamoDB for storage.",
"Set Up AWS Account",
"Create Lambda Function",
"APIGateway Integration",
"Define DynamoDB Table Scheme",
"Connect Service To DB",
"Add Authentication Layers",
"Monitor Metrics and Set Alarms"],

["Migrate our current monolithic PHP application towards containerized services using Docker and Kubernetes for scalability.",
"Architectural Restructuring Plan",
"Containerisation Process With Docker",
"Service Orchestration With Kubernetes",
"Load Balancing Strategies",
"Persistent Storage Solutions",
"Network Policies Enforcement",
"Continuous Integration / Continuous Delivery"],

["Provide guidance on integrating WebAssembly modules compiled from C++ source files into an ongoing web project.",
"Toolchain Selection (Emscripten vs. LLVM)",
"Setting Up Compiler Environment",
".cpp Source Preparation",
"Module Building Approach",
"Memory Management Considerations",
"Performance Tradeoffs",
"Seamless Web Assembly Embedding"]
]

def parse_action(line):
    action_name, action_input = line.strip("action: ").split("=")
    action_input = action_input.strip()
    return action_name, action_input

def get_file_tree(path):
    """
    Recursively explores a directory and returns a nested dictionary representing its file tree.
    """
    tree = {}
    for item in os.listdir(path):
        item_path = os.path.join(path, item)
        if os.path.isdir(item_path):
            tree[item] = get_file_tree(item_path)
        else:
            tree[item] = None
    return tree

def display_file_tree(tree, indent=0):
    """
    Prints a formatted representation of the file tree.
    """
    for name, subtree in tree.items():
        print(f"{'  ' * indent}{name}")
        if subtree is not None:
            display_file_tree(subtree, indent + 1)

def project_explorer(path):
    """
    Displays the file tree of a given path in a Streamlit app.
    """
    tree = get_file_tree(path)
    display_file_tree(tree)

def chat_app_logic(message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
    # Your existing code here
    
    try:
        # Attempt to join the generator output
        response = ''.join(generate(
            model=model,
            messages=messages,
            stream=True,
            temperature=0.7,
            max_tokens=1500
        ))
    except TypeError:
        # If joining fails, collect the output in a list
        response_parts = []
        for part in generate(
            model=model,
            messages=messages,
            stream=True,
            temperature=0.7,
            max_tokens=1500
        ):
            if isinstance(part, str):
                response_parts.append(part)
            elif isinstance(part, dict) and 'content' in part:
                response_parts.append(part['content']),
        
        response = ''.join(response_parts,
        # Run the model and get the response (convert generator to string)
            prompt=message,
            history=history,
            agent_name=agent_name,
            sys_prompt=sys_prompt,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
                          )
        history.append((message, response))
        return history

    return history

def main():
    with gr.Blocks() as demo:
        gr.Markdown("## FragMixt")
        gr.Markdown("### Agents w/ Agents")

        # Chat Interface
        chatbot = gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel")
        #chatbot.load(examples)

        # Input Components
        message = gr.Textbox(label="Enter your message", placeholder="Ask me anything!")
        purpose = gr.Textbox(label="Purpose", placeholder="What is the purpose of this interaction?")
        agent_name = gr.Dropdown(label="Agents", choices=[s for s in agents], value=agents[0], interactive=True)
        sys_prompt = gr.Textbox(label="System Prompt", max_lines=1, interactive=True)
        temperature = gr.Slider(label="Temperature", value=0.9, minimum=0.0, maximum=1.0, step=0.05, interactive=True, info="Higher values produce more diverse outputs")
        max_new_tokens = gr.Slider(label="Max new tokens", value=1048*10, minimum=0, maximum=1048*10, step=64, interactive=True, info="The maximum numbers of new tokens")
        top_p = gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.0, maximum=1, step=0.05, interactive=True, info="Higher values sample more low-probability tokens")
        repetition_penalty = gr.Slider(label="Repetition penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05, interactive=True, info="Penalize repeated tokens")

        # Button to submit the message
        submit_button = gr.Button(value="Send")

        # Project Explorer Tab
        with gr.Tab("Project Explorer"):
            project_path = gr.Textbox(label="Project Path", placeholder="/home/user/app/current_project")
            explore_button = gr.Button(value="Explore")
            project_output = gr.Textbox(label="File Tree", lines=20)

        # Chat App Logic Tab
        with gr.Tab("Chat App"):
            history = gr.State([])
            for example in examples:
                gr.Button(value=example[0]).click(lambda: chat_app_logic(example[0], history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty), outputs=chatbot)

            # Connect components to the chat app logic
            submit_button.click(chat_app_logic, inputs=[message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty], outputs=chatbot)
            message.submit(chat_app_logic, inputs=[message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty], outputs=chatbot)

        # Connect components to the project explorer
        explore_button.click(project_explorer, inputs=project_path, outputs=project_output)

    demo.launch(show_api=True)

if __name__ == "__main__":
    main()