128 lines
4.8 KiB
Python
128 lines
4.8 KiB
Python
import os
|
|
import json
|
|
import requests
|
|
from flask import Blueprint, request, jsonify
|
|
from flask_login import login_required
|
|
|
|
llm = Blueprint('llm', __name__)
|
|
|
|
def _call_llm(action, natural_query, code, runtime, logs):
|
|
api_key = os.environ.get("GEMINI_API_KEY")
|
|
model = os.environ.get("GEMINI_MODEL", "gemini-2.0-flash")
|
|
|
|
if not api_key:
|
|
return None, "GEMINI_API_KEY not set."
|
|
|
|
api_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
|
|
headers = {'Content-Type': 'application/json'}
|
|
|
|
system_instruction = ""
|
|
user_prompt = ""
|
|
|
|
if action == "generate":
|
|
system_instruction = f"""
|
|
You are an expert {runtime} developer. Your task is to write a complete, production-ready {runtime} async function based on the user's request.
|
|
|
|
**Function Signature:**
|
|
Your function MUST have the following signature: `async (req, environment) => {{ ... }}` (for Node/Deno) or `def main(req, environment):` (for Python).
|
|
- `req`: An object containing details about the incoming HTTP request.
|
|
- `environment`: A mutable JSON object that persists across executions.
|
|
|
|
**Constraints:**
|
|
- The function will be executed in a sandboxed environment.
|
|
- Define ALL helper functions inside the main function.
|
|
- DO NOT use external module loading (require/import) unless standard lib. DO NOT IMPORT!
|
|
- Return ONLY the code.
|
|
"""
|
|
user_prompt = f"Write a function that: {natural_query}"
|
|
|
|
elif action == "explain":
|
|
system_instruction = f"You are an expert {runtime} developer. Explain what the following code does in simple, concise terms suitable for a developer."
|
|
user_prompt = f"Code:\n```\n{code}\n```"
|
|
|
|
elif action == "optimize":
|
|
system_instruction = f"""
|
|
You are an expert {runtime} developer. Analyze the following code and provide an optimized version.
|
|
Focus on performance, readability, and best practices.
|
|
Return ONLY the optimized code, without markdown formatting or explanation.
|
|
"""
|
|
user_prompt = f"Code:\n```\n{code}\n```"
|
|
|
|
elif action == "debug":
|
|
system_instruction = f"""
|
|
You are an expert {runtime} developer. Analyze the following code and error logs.
|
|
Explain the error and provide a fixed version of the code.
|
|
Format your response as:
|
|
**Explanation:** [Brief explanation of the bug]
|
|
**Fix:**
|
|
```
|
|
[Fixed Code]
|
|
```
|
|
"""
|
|
user_prompt = f"Code:\n```\n{code}\n```\n\nLogs:\n{logs}"
|
|
|
|
try:
|
|
payload = json.dumps({
|
|
"contents": [{"parts": [{"text": system_instruction + "\n\n" + user_prompt}]}]
|
|
})
|
|
|
|
response = requests.post(api_url, headers=headers, data=payload)
|
|
response.raise_for_status()
|
|
|
|
response_data = response.json()
|
|
|
|
candidates = response_data.get('candidates', [])
|
|
if not candidates:
|
|
return None, "No candidates found in API response."
|
|
|
|
content = candidates[0].get('content', {})
|
|
parts = content.get('parts', [])
|
|
if not parts:
|
|
return None, "No parts found in API response content."
|
|
|
|
generated_text = parts[0].get('text', '').strip()
|
|
|
|
# Clean up code blocks for generate/optimize actions
|
|
if action in ["generate", "optimize"]:
|
|
if generated_text.startswith("```"):
|
|
# Find first newline
|
|
first_newline = generated_text.find("\n")
|
|
if first_newline != -1:
|
|
generated_text = generated_text[first_newline+1:]
|
|
if generated_text.endswith("```"):
|
|
generated_text = generated_text[:-3]
|
|
generated_text = generated_text.strip()
|
|
|
|
return generated_text, None
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
return None, f"Error communicating with API: {e}"
|
|
except (KeyError, IndexError, Exception) as e:
|
|
return None, f"Error processing API response: {e}"
|
|
|
|
@llm.route("/generate_script", methods=["POST"])
|
|
@login_required
|
|
def generate_script():
|
|
try:
|
|
data = request.json
|
|
action = data.get('action', 'generate') # Default to generate for backward compatibility
|
|
natural_query = data.get('natural_query')
|
|
code = data.get('code')
|
|
runtime = data.get('runtime', 'node')
|
|
logs = data.get('logs')
|
|
|
|
if action == 'generate' and not natural_query:
|
|
return jsonify({"error": "natural_query is required for generation"}), 400
|
|
|
|
if action in ['explain', 'optimize', 'debug'] and not code:
|
|
return jsonify({"error": "code is required for this action"}), 400
|
|
|
|
result, error = _call_llm(action, natural_query, code, runtime, logs)
|
|
|
|
if error:
|
|
return jsonify({"error": error}), 500
|
|
|
|
return jsonify({"result": result})
|
|
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500 |