Update codex/code_generator.py

codex/code_generator.py

@@ -5,6 +5,9 @@ HF_API = "https://api-inference.huggingface.co/models/bigcode/starcoder"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 def generate_code(task):
+    if not HF_TOKEN:
+        raise RuntimeError("HF_TOKEN not set")
+
     headers = {"Authorization": f"Bearer {HF_TOKEN}"}
     prompt = f"""
 Write safe Python code to accomplish this task:
@@ -15,11 +18,15 @@ Rules:
 - No OS commands
 - No file deletion
 - No network calls
-- No infinite loops
-- Return only code
+- Must define a function
+- Must assign final output to variable named `result`
+- Return only Python code
 """
 
-    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 400}}
+    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 300}}
     r = requests.post(HF_API, headers=headers, json=payload, timeout=60)
+    r.raise_for_status()
+
     text = r.json()[0]["generated_text"]
-    return text.split("```")[-1]
+    code = text.split("```")[-1].strip()
+    return code