agents/ui_designer_agent.py

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from langchain_core.messages import AIMessage

MODEL_REPO = "Rahul-8799/product_manager_mistral"

tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_REPO,
    torch_dtype=torch.float16,
    device_map="auto"
)

def run(state: dict) -> dict:
    """UI Designer creates beautiful and structured UI designs with proper spacing and layout"""
    messages = state["messages"]
    prompt = messages[-1].content
    
    # Enhance the prompt with UI design principles
    enhanced_prompt = f"""
    Create a beautiful and well-structured UI design following these principles:
    1. Use proper spacing and padding (recommended: 1rem/16px for padding, 2rem/32px for margins)
    2. Implement a consistent color scheme
    3. Ensure proper hierarchy with clear headings
    4. Use responsive design principles
    5. Implement proper grid system
    6. Add smooth transitions and hover effects
    7. Ensure proper contrast and readability
    8. Use modern UI components and patterns
    
    Original requirements: {prompt}
    """
    
    input_ids = tokenizer(enhanced_prompt, return_tensors="pt").input_ids.to(model.device)
    output_ids = model.generate(input_ids, max_new_tokens=3000)
    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    return {
        "messages": [AIMessage(content=output)],
        "chat_log": state["chat_log"] + [{"role": "UI Designer", "content": output}],
        "ui_design_output": output,
    }