You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
76 lines
2.2 KiB
76 lines
2.2 KiB
3 months ago
|
from flask import Blueprint, jsonify, request
|
||
|
|
import anthropic
|
||
|
|
import os
|
||
|
|
from openai import OpenAI
|
||
|
|
|
||
|
|
code_generation_bp = Blueprint('code_generation', __name__)
|
||
|
|
|
||
|
|
@code_generation_bp.route('/v1/code/generation', methods=['POST'])
|
||
|
|
def generate_code():
|
||
|
|
body = request.json
|
||
|
|
prompt_content = body.get("prompt", "")
|
||
|
|
model = body.get("model", "gpt-4o")
|
||
|
|
|
||
|
|
general_system_prompt = """
|
||
|
|
You are an assistant that knows about Manim. Manim is a mathematical animation engine that is used to create videos programmatically.
|
||
|
|
|
||
|
|
The following is an example of the code:
|
||
|
|
\`\`\`
|
||
|
|
from manim import *
|
||
|
|
from math import *
|
||
|
|
|
||
|
|
class GenScene(Scene):
|
||
|
|
def construct(self):
|
||
|
|
c = Circle(color=BLUE)
|
||
|
|
self.play(Create(c))
|
||
|
|
|
||
|
|
\`\`\`
|
||
|
|
|
||
|
|
# Rules
|
||
|
|
1. Always use GenScene as the class name, otherwise, the code will not work.
|
||
|
|
2. Always use self.play() to play the animation, otherwise, the code will not work.
|
||
|
|
3. Do not use text to explain the code, only the code.
|
||
|
|
4. Do not explain the code, only the code.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if model.startswith("claude-"):
|
||
|
|
client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
|
||
|
|
messages = [{"role": "user", "content": prompt_content}]
|
||
|
|
try:
|
||
|
|
response = client.messages.create(
|
||
|
|
model=model,
|
||
|
|
max_tokens=1000,
|
||
|
|
temperature=0.2,
|
||
|
|
system=general_system_prompt,
|
||
|
|
messages=messages,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Extract the text content from the response
|
||
|
|
code = "".join(block.text for block in response.content)
|
||
|
|
|
||
|
|
return jsonify({"code": code})
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
return jsonify({"error": str(e)}), 500
|
||
|
|
|
||
|
|
else:
|
||
|
|
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
|
||
|
|
messages = [
|
||
|
|
{"role": "system", "content": general_system_prompt},
|
||
|
|
{"role": "user", "content": prompt_content},
|
||
|
|
]
|
||
|
|
|
||
|
|
try:
|
||
|
|
response = client.chat.completions.create(
|
||
|
|
model=model,
|
||
|
|
messages=messages,
|
||
|
|
temperature=0.2,
|
||
|
|
)
|
||
|
|
|
||
|
|
code = response.choices[0].message.content
|
||
|
|
|
||
|
|
return jsonify({"code": code})
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
return jsonify({"error": str(e)}), 500
|