Update backend_agent/api_generator.py

backend_agent/api_generator.py

@@ -9,25 +9,27 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 def generate_backend_code_llm(task_name):
     """
-    Generates a Python backend skeleton (Flask API) for a given backend task
-    using an LLM.
+    Generates Python backend code (REST or GraphQL) for a given task using an LLM.
     
-    Returns a string containing the code.
+    The LLM decides whether the task is best suited for a REST API or GraphQL API
+    based on the task description.
     """
     prompt = f"""
-You are an expert backend developer. Generate a Python Flask API skeleton for this backend task.
-Include routes, basic input validation, and placeholders for business logic.
+You are an expert backend developer. Given this task, generate Python code
+for either a REST API (Flask/FastAPI) or a GraphQL API (Graphene/FastAPI-GraphQL),
+depending on which is more appropriate for the task.
 
 Task: {task_name}
 
-Python Flask code:
+Provide the Python code with routes, input validation, and placeholders for business logic.
+Also include comments explaining what each part does.
 """
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=250)
+    outputs = model.generate(**inputs, max_new_tokens=350)
     code = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Clean up the prompt portion from the output
-    if "Python Flask code:" in code:
-        code = code.split("Python Flask code:")[-1].strip()
+    # Remove prompt from output if it’s repeated
+    if "Task:" in code:
+        code = code.split("Task:")[-1].strip()
 
     return code