Update frontend_agent/ui_generator.py

frontend_agent/ui_generator.py

@@ -8,17 +8,18 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 def generate_react_component_llm(task_name):
     """
-    Generates a React functional component for a frontend task
+    Uses LLM to generate a React functional component for a frontend task.
     """
     prompt = f"""
-Generate a React functional component for this frontend task:
+Generate a React functional component for this frontend task.
 Task: {task_name}
-Include basic JSX structure and comments explaining each part.
+Include props if needed and simple placeholder JSX.
 """
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=250)
     code = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    if "Component:" in code:
-        code = code.split("Component:")[-1].strip()
+    # Remove prompt part if included
+    if "Task:" in code:
+        code = code.split("Task:")[-1].strip()
     return code