简介

本文介绍如何使用 Ollama API 来自动生成 Python 代码。我们将以生成一个计算数字列表总和的函数为例,展示整个实现过程。

实现步骤

1. 环境准备

首先需要导入必要的Python模块:

import requests  # 用于发送HTTP请求
import json     # 用于处理JSON数据
import subprocess  # 用于执行生成的代码

2. 配置API请求

设置 Ollama API 的访问地址和请求参数:

api_url = "http://localhost:11434/api/generate"

data = {
    "model": "codegemma:2b",  # 使用的AI模型
    "prompt": "编写一个Python函数来计算数字列表的总和。\ndef calculate_sum(numbers):\n# 调用函数并打印结果\n# 示例: print(calculate_sum([1, 2, 3, 4, 5]))",
    "options": {
        "temperature": 0.2,    # 控制生成代码的创造性
        "max_tokens": 100      # 限制生成代码的长度
    }
}

3. 发送API请求

向 Ollama API 发送POST请求并获取响应:

response = requests.post(
    api_url, 
    headers={"Content-Type": "application/json"}, 
    data=json.dumps(data), 
    stream=True
)

4. 处理响应结果

如果请求成功,我们需要:

  1. 解析返回的数据流
  2. 收集生成的代码
  3. 将代码保存到文件
  4. 执行生成的代码
if response.status_code == 200:
    generated_code = ""
    
    # 逐行处理响应内容
    for line in response.iter_lines():
        if line:
            line_data = json.loads(line.decode('utf-8'))
            generated_code += line_data.get("response", "")
            if line_data.get("done", False):
                break
                
    # 保存代码到文件
    with open("gensum.py", "w") as code_file:
        code_file.write(generated_code)
        
    # 执行生成的代码
    result = subprocess.run(["python", "gensum.py"], capture_output=True, text=True)
    print(result.stdout)

代码特点

  1. 流式处理:使用流式响应处理大量数据
  2. 错误处理:包含基本的状态码检查
  3. 文件操作:自动将生成的代码保存到文件
  4. 代码执行:自动运行生成的代码并显示结果

使用说明

  1. 确保本地已启动 Ollama 服务
  2. 运行代码后,会在同目录下生成 gensum.py 文件
  3. 程序会自动执行生成的代码并显示计算结果

注意事项

  • 确保本地 Ollama 服务正常运行
  • 检查网络连接是否正常
  • 确保有足够的磁盘空间保存生成的代码文件
  • `import ollama
    import json

定义生成代码的函数

def generate_code():
try:
# 配置提示词
prompt = “”"编写一个Python函数来计算数字列表的总和。
def calculate_sum(numbers):

调用函数并打印结果

示例: print(calculate_sum([1, 2, 3, 4, 5]))“”"

    # 使用 Ollama 生成代码
    response = ollama.generate(
        model='codegemma:2b',
        prompt=prompt,
        options={
            "temperature": 0.2,
            "num_predict": 100
        }
    )

    # 获取生成的代码
    generated_code = response['response']
    print("生成的代码:\n", generated_code)

    # 将代码保存到文件
    with open("gensumo.py", "w") as code_file:
        code_file.write(generated_code)

    # 执行生成的代码
    print("\n执行生成的代码:")
    exec(generated_code)

except Exception as e:
    print(f"发生错误: {str(e)}")

if name == “main”:
generate_code()`

这个实现展示了如何利用 AI 模型来自动生成和执行 Python 代码,可以作为代码生成自动化的基础框架。

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐