Update backend_agent/api_generator.py

backend_agent/api_generator.py

@@ -0,0 +1,33 @@
+# backend_agent/api_generator.py
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# Use a decoder-only LLM (OPT-125m for demo, can scale up)
+MODEL_NAME = "facebook/opt-125m"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+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.
+    
+    Returns a string containing the code.
+    """
+    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.
+
+Task: {task_name}
+
+Python Flask code:
+"""
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=250)
+    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()
+
+    return code