import os
import openai
from dotenv import load_dotenv
load_dotenv()
class CodeAssistant:
def __init__(self):
"""初始化代码助手"""
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise ValueError("请设置OPENAI_API_KEY环境变量")
openai.api_key = api_key
self.model = "gpt-3.5-turbo"
def generate(self, prompt, language="python"):
"""
根据描述生成代码
参数:
prompt: 自然语言描述
language: 目标语言
返回:
生成的代码
"""
system_msg = f"""你是一个{language}专家。请根据需求生成完整、可运行的代码。
要求:
1. 只返回代码,不要解释
2. 代码应包含必要注释
3. 确保代码可以直接运行
需求:{prompt}"""
try:
response = openai.ChatCompletion.create(
model=self.model,
messages=[
{"role": "system", "content": system_msg},
{"role": "user", "content": prompt}
],
temperature=0.7,
max_tokens=800
)
code = response.choices[0].message.content.strip()
return self._clean_code(code)
except Exception as e:
return f"生成失败:{str(e)}"
def _clean_code(self, code):
"""清理代码格式"""
if code.startswith("```"):
lines = code.split("\n")
code = "\n".join(lines[1:-1]) if lines[-1].startswith("```") else "\n".join(lines[1:])
return code
def save(self, code, filename="output.py"):
"""保存代码到文件"""
with open(filename, "w", encoding="utf-8") as f:
f.write(code)
print(f"✅ 已保存:{filename}")
if __name__ == "__main__":
assistant = CodeAssistant()
print("🤖 智能代码助手")
print("输入需求或'quit'退出\n")
while True:
user_input = input("> ")
if user_input.lower() in ['quit', 'exit']:
break
if not user_input.strip():
continue
print("生成中...")
code = assistant.generate(user_input)
print("\n" + "="*40)
print(code)
print("="*40)
if input("保存?(y/n): ").lower() == 'y':
name = input("文件名: ") or "output.py"
assistant.save(code, name)