通过大模型执行 Python 代码

package nexus.io.gemini;

import java.util.HashMap;
import java.util.Map;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;

@Data
@NoArgsConstructor
@AllArgsConstructor
public class GeminiResponseSchema {
  private String type = "object";
  private Map<String, GeminiResponseProperty> properties;

  /**
    "responseSchema": {
      "type": "object",
      "properties": {
        "suggestions": {
          "type": "array",
          "items": {
            "type": "string"
          }
        }
      }
    }
   */
  public static GeminiResponseSchema array(String key) {
    Map<String, String> items = itemType("string");
    Map<String, GeminiResponseProperty> properties = new HashMap<>();
    properties.put(key, new GeminiResponseProperty("array", items));

    GeminiResponseSchema schema = new GeminiResponseSchema();
    schema.setProperties(properties);
    return schema;
  }

  /**
   *
    "responseSchema": {
      "type": "object",
      "properties": {
        "action": {
          "type": "string"
        },
        "tool": {
          "type": "string"
        }
        "code": {
          "type": "string"
        },

      }
    }
   * @return
   */
  public static GeminiResponseSchema pythonCode() {
    Map<String, GeminiResponseProperty> properties = new HashMap<>();
    properties.put("action", new GeminiResponseProperty("string"));
    properties.put("tool", new GeminiResponseProperty("string"));
    properties.put("code", new GeminiResponseProperty("string"));
    GeminiResponseSchema schema = new GeminiResponseSchema();
    schema.setProperties(properties);
    return schema;
  }

  private static Map<String, String> itemType(String value) {
    Map<String, String> items = new HashMap<>();
    items.put("type", value);
    return items;
  }
}

package nexus.io.llm.service;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.List;

import nexus.io.gemini.GeminiCandidateVo;
import nexus.io.gemini.GeminiChatRequestVo;
import nexus.io.gemini.GeminiChatResponseVo;
import nexus.io.gemini.GeminiClient;
import nexus.io.gemini.GeminiGenerationConfigVo;
import nexus.io.gemini.GeminiPartVo;
import nexus.io.gemini.GeminiResponseSchema;
import nexus.io.gemini.GoogleGeminiModels;
import nexus.io.jfinal.aop.Aop;
import nexus.io.linux.LinuxClient;
import nexus.io.linux.ProcessResult;
import nexus.io.llm.vo.ToolVo;
import nexus.io.template.PromptEngine;
import nexus.io.tio.boot.admin.services.AwsS3StorageService;
import nexus.io.model.upload.UploadResult;
import nexus.io.model.upload.UploadFile;
import nexus.io.tio.utils.encoder.Base64Utils;
import nexus.io.tio.utils.encoder.ImageVo;
import nexus.io.tio.utils.hutool.StrUtil;
import nexus.io.tio.utils.json.JsonUtils;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MatplotlibService {

  public ProcessResult generateMatplot(String quesiton, String answer) {
    String text = generateCode(quesiton, answer);
    if (StrUtil.isBlank(text)) {
      return null;
    }
    ToolVo toolVo = null;
    try {
      toolVo = JsonUtils.parse(text, ToolVo.class);
    } catch (Exception e) {
      log.error("text:{}", text, e.getMessage(), e);
      StringWriter stringWriter = new StringWriter();
      PrintWriter printWriter = new PrintWriter(stringWriter);
      e.printStackTrace(printWriter);
      String stackTrace = stringWriter.toString();
      String msg = "code:" + text + ",stackTrace" + stackTrace;
      Aop.get(AgentNotificationService.class).sendError(msg);
      return null;
    }
    if (toolVo != null && "execute_python".equals(toolVo.getTool())) {
      String code = toolVo.getCode();
      ProcessResult result = LinuxClient.executePythonCode(code);
      if (result != null) {
        String stdErr = result.getStdErr();
        if (StrUtil.isNotBlank(stdErr)) {
          String prompt = "python代码执行过程中出现了错误,请修正错误并仅输出修改后的代码,错误信息:%s";
          prompt = String.format(prompt, stdErr);
          code = fixCodeError(prompt, code);
          result = LinuxClient.executePythonCode(code);
        }

        List<String> images = result.getImages();
        if (images != null) {
          List<String> imageUrls = new ArrayList<>(images.size());
          for (String imageBase64Code : images) {
            ImageVo decodeImage = Base64Utils.decodeImage(imageBase64Code);
            UploadFile uploadFile = new UploadFile("matplotlib." + decodeImage.getExtension(), decodeImage.getData());
            UploadResult resultVo = Aop.get(AwsS3StorageService.class).uploadFile("matplotlib", uploadFile);
            String url = resultVo.getUrl();
            imageUrls.add(url);
          }
          result.setImages(imageUrls);
        }
        return result;
      }

    }
    return null;
  }

  private String fixCodeError(String userPrompt, String code) {
    String text = generateCode(userPrompt, code);
    if (StrUtil.isBlank(text)) {
      return null;
    }
    ToolVo toolVo = null;
    try {
      toolVo = JsonUtils.parse(text, ToolVo.class);
      return toolVo.getCode();
    } catch (Exception e) {
      log.error("text:{}", text, e.getMessage(), e);
      StringWriter stringWriter = new StringWriter();
      PrintWriter printWriter = new PrintWriter(stringWriter);
      e.printStackTrace(printWriter);
      String stackTrace = stringWriter.toString();
      String msg = "code:" + text + ",stackTrace" + stackTrace;
      Aop.get(AgentNotificationService.class).sendError(msg);
      return null;
    }
  }

  public String generateCode(String quesiton, String answer) {
    String systemPrompt = PromptEngine.renderToString("python_code_prompt.txt");
    String userPrompt = "请根据下面的用户的问题和答案使用python代码绘制函数图像帮助用户更好的理解问题.如果无需绘制函数图像,请返回`not_needed`";

    // 1. Construct request body
    GeminiChatRequestVo reqVo = new GeminiChatRequestVo();
    reqVo.setUserPrompts(userPrompt, quesiton, answer);
    reqVo.setSystemPrompt(systemPrompt);

    GeminiResponseSchema pythonCode = GeminiResponseSchema.pythonCode();
    GeminiGenerationConfigVo geminiGenerationConfigVo = new GeminiGenerationConfigVo();
    geminiGenerationConfigVo.buildJsonValue().setResponseSchema(pythonCode);

    reqVo.setGenerationConfig(geminiGenerationConfigVo);

    // 2. Send sync request: generateContent
    GeminiChatResponseVo respVo = GeminiClient.generate(GoogleGeminiModels.GEMINI_2_0_FLASH, reqVo);
    if (respVo != null) {
      List<GeminiCandidateVo> candidates = respVo.getCandidates();
      GeminiCandidateVo candidate = candidates.get(0);
      List<GeminiPartVo> parts = candidate.getContent().getParts();
      if (candidate != null && candidate.getContent() != null && parts != null) {
        String text = parts.get(0).getText();
        return text;
      }
    }
    return null;
  }
}

package nexus.io.linux;

import java.io.IOException;

import nexus.io.tio.utils.environment.EnvUtils;
import nexus.io.tio.utils.http.OkHttpClientPool;
import nexus.io.tio.utils.json.JsonUtils;

import okhttp3.Call;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;

public class LinuxClient {

  public static ProcessResult executePythonCode(String code) {
    String apiBase = EnvUtils.getStr("LINUX_BASE_URL");
    String key = EnvUtils.getStr("LINUX_API_KEY");

    String arg0 = apiBase + "/python";
    MediaType mediaType = MediaType.parse("text/plain");

    OkHttpClient client = OkHttpClientPool.get60HttpClient();
    RequestBody body = RequestBody.create(code, mediaType);
    Request request = new Request.Builder().url(arg0).method("POST", body).addHeader("authorization", "Bearer " + key).build();
    Call call = client.newCall(request);
    try (Response response = call.execute()) {
      String string = response.body().string();
      int resposneCode = response.code();
      if (response.isSuccessful()) {
        return JsonUtils.parse(string, ProcessResult.class);
      } else {
        throw new RuntimeException("code:" + resposneCode + " response:" + string);
      }
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }
}

package llm;

import java.net.URL;

import org.junit.Test;

import nexus.io.jfinal.aop.Aop;
import nexus.io.linux.ProcessResult;
import nexus.io.llm.service.MatplotlibService;
import nexus.io.tio.utils.environment.EnvUtils;
import nexus.io.tio.utils.hutool.FileUtil;
import nexus.io.tio.utils.hutool.ResourceUtil;

public class MatplotlibServiceTest {

  @Test
  public void test() {
    EnvUtils.load();
    URL resource = ResourceUtil.getResource("data/example_answer.txt");
    if (resource == null) {
      return;
    }
    StringBuilder readString = FileUtil.readString(resource);
    String generateCode = Aop.get(MatplotlibService.class).generateCode("解题",readString.toString());
//    String generateCode = Aop.get(MatplotlibService.class).generateCode("解题","我爱");
    System.out.println(generateCode);
  }

  @Test
  public void generateMatplot() {
    EnvUtils.load();

    URL resource = ResourceUtil.getResource("data/example_answer.txt");
    if (resource == null) {
      return;
    }
    StringBuilder readString = FileUtil.readString(resource);

    for (int i = 0; i < 10; i++) {
      try {
        ProcessResult generateMatplot = Aop.get(MatplotlibService.class).generateMatplot("解题", readString.toString());
        System.out.println(generateMatplot);
      } catch (Exception e) {
        e.printStackTrace();
      }
    }
  }
}