app.py

import os
import gradio as gr
import torch
import json
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# Set Hugging Face Token for Authentication (ensure it's set in your environment)
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

# Base model (needed for QLoRA adapter)
BASE_MODEL = "meta-llama/Llama-3-1B-Instruct"
QLORA_ADAPTER = "meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8"

# Function to load Llama model
def load_llama_model():
    print("Loading base model...")
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16 if torch.has_bfloat16 else torch.float32,  # Use bfloat16 if available, else float32
        device_map="cpu",  # Ensure it runs on CPU
        token=HUGGINGFACE_TOKEN
    )

    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False, token=HUGGINGFACE_TOKEN)

    print("Loading QLoRA adapter...")
    model = PeftModel.from_pretrained(
        model,
        QLORA_ADAPTER,
        token=HUGGINGFACE_TOKEN
    )
    
    print("Merging LoRA weights...")
    model = model.merge_and_unload()  # Merge LoRA weights for inference

    return tokenizer, model

# Load Llama 3.2 model
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8"
tokenizer, model = load_llama_model(MODEL_NAME)

# Load Llama Guard for content moderation
LLAMA_GUARD_NAME = "meta-llama/Llama-Guard-3-1B-INT4"
guard_tokenizer, guard_model = load_llama_model(LLAMA_GUARD_NAME)

# Define Prompt Templates
PROMPTS = {
    "project_analysis": """<|begin_of_text|><|prompt|>Analyze this project description and generate:
1. Project timeline with milestones
2. Required technology stack
3. Potential risks
4. Team composition
5. Cost estimation

Project: {project_description}<|completion|>""",
    
    "code_generation": """<|begin_of_text|><|prompt|>Generate implementation code for this feature:
{feature_description}

Considerations:
- Use {programming_language}
- Follow {coding_standards}
- Include error handling
- Add documentation<|completion|>""",

    "risk_analysis": """<|begin_of_text|><|prompt|>Predict potential risks for this project plan:
{project_data}

Format output as JSON with risk types, probabilities, and mitigation strategies<|completion|>"""
}

# Function: Content Moderation using Llama Guard
def moderate_input(user_input):
    # Llama Guard specific prompt format
    prompt = f"""<|begin_of_text|><|user|>
Input: {user_input}
Please verify that this input doesn't violate any content policies.
<|assistant|>"""
    
    inputs = guard_tokenizer(prompt, return_tensors="pt", truncation=True)
    
    with torch.no_grad():  # Disable gradient calculation for inference
        outputs = guard_model.generate(
            inputs.input_ids,
            max_length=256,
            temperature=0.1
        )
    
    response = guard_tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if "flagged" in response.lower() or "violated" in response.lower() or "policy violation" in response.lower():
        return "⚠️ Content flagged by Llama Guard. Please modify your input."
    return None  # Safe input, proceed normally

# Function: Generate AI responses
def generate_response(prompt_type, **kwargs):
    prompt = PROMPTS[prompt_type].format(**kwargs)
    
    moderation_warning = moderate_input(prompt)
    if moderation_warning:
        return moderation_warning  # Stop processing if flagged

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
    
    with torch.no_grad():  # Disable gradient calculation for inference
        outputs = model.generate(
            inputs.input_ids,
            max_length=1024,
            temperature=0.7 if prompt_type == "project_analysis" else 0.5,
            top_p=0.9,
            do_sample=True
        )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Function: Analyze project
def analyze_project(project_desc):
    return generate_response("project_analysis", project_description=project_desc)

# Function: Generate code
def generate_code(feature_desc, lang="Python", standards="PEP8"):
    return generate_response("code_generation", feature_description=feature_desc, programming_language=lang, coding_standards=standards)

# Function: Predict risks
def predict_risks(project_data):
    risks = generate_response("risk_analysis", project_data=project_data)
    try:
        # Try to extract JSON part from the response
        import re
        json_match = re.search(r'\{.*\}', risks, re.DOTALL)
        if json_match:
            return json.loads(json_match.group(0))
        return {"error": "Could not parse JSON response"}
    except json.JSONDecodeError:
        return {"error": "Invalid JSON response. Please refine your input."}

# Gradio UI
def create_gradio_interface():
    with gr.Blocks(title="AI Project Manager", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🚀 AI-Powered Project Manager & Code Assistant")
        
        # Project Analysis Tab
        with gr.Tab("Project Setup"):
            project_input = gr.Textbox(label="Project Description", lines=5, placeholder="Describe your project...")
            project_output = gr.Textbox(label="Project Analysis", lines=15)  # Changed from JSON to Textbox for better formatting
            analyze_btn = gr.Button("Analyze Project")
            analyze_btn.click(analyze_project, inputs=project_input, outputs=project_output)
        
        # Code Generation Tab
        with gr.Tab("Code Assistant"):
            code_input = gr.Textbox(label="Feature Description", lines=3)
            lang_select = gr.Dropdown(["Python", "JavaScript", "Java", "C++"], label="Language", value="Python")
            standards_select = gr.Dropdown(["PEP8", "Google", "Airbnb"], label="Coding Standard", value="PEP8")
            code_output = gr.Code(label="Generated Code")
            code_btn = gr.Button("Generate Code")
            code_btn.click(generate_code, inputs=[code_input, lang_select, standards_select], outputs=code_output)
        
        # Risk Analysis Tab
        with gr.Tab("Risk Analysis"):
            risk_input = gr.Textbox(label="Project Plan", lines=5)
            risk_output = gr.JSON(label="Risk Predictions") 
            risk_btn = gr.Button("Predict Risks")
            risk_btn.click(predict_risks, inputs=risk_input, outputs=risk_output)
        
        # Real-time Chatbot for Collaboration
        with gr.Tab("Live Collaboration"):
            gr.Markdown("## Real-time Project Collaboration")
            chat = gr.Chatbot(height=400)
            msg = gr.Textbox(label="Chat with AI PM")
            clear = gr.Button("Clear Chat")
            
            def respond(message, chat_history):
                moderation_warning = moderate_input(message)
                if moderation_warning:
                    chat_history.append((message, moderation_warning))
                    return "", chat_history

                # Format chat history for context
                history_text = ""
                for i, (usr, ai) in enumerate(chat_history[-3:]):  # Use last 3 messages for context
                    history_text += f"User: {usr}\nAI: {ai}\n"
                
                prompt = f"""<|begin_of_text|><|prompt|>Project Management Chat:
Context: {message}
Chat History: {history_text}
User: {message}<|completion|>"""
                
                inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
                
                with torch.no_grad():
                    outputs = model.generate(
                        inputs.input_ids,
                        max_length=1024,
                        temperature=0.7,
                        top_p=0.9,
                        do_sample=True
                    )
                
                response = tokenizer.decode(outputs[0], skip_special_tokens=True)
                chat_history.append((message, response))
                return "", chat_history
            
            msg.submit(respond, [msg, chat], [msg, chat])
            clear.click(lambda: None, None, chat, queue=False)

    return demo

# Run Gradio App
if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch(share=True)