Quellcode durchsuchen

ai 添加调用ragflow

yangshun vor 3 Wochen
Ursprung
Commit
8ff18d1a2c

+ 122 - 4
guoyan-ai/src/main/java/com/cy/guoyan/admin/module/ai/util/DifyClient.java

@@ -1,6 +1,7 @@
 package com.cy.guoyan.admin.module.ai.util;
 
 import com.alibaba.fastjson.JSONObject;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.http.client.methods.*;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
@@ -12,10 +13,12 @@ import org.apache.http.util.EntityUtils;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
 import javax.annotation.PostConstruct;
-import java.io.File;
-import java.io.IOException;
+import java.io.*;
+import java.net.HttpURLConnection;
+import java.net.URL;
 import java.net.URLEncoder;
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -23,6 +26,7 @@ import java.util.Map;
  * Dify AI HTTP 客户端工具类
  */
 @Component
+@Slf4j
 public class DifyClient  {
 
     private static final String AUTH_HEADER       = "Authorization";
@@ -30,10 +34,16 @@ public class DifyClient  {
     private static final String BEARER            = "Bearer ";
 
     @Value("${dify.baseUrl}")
-    private String baseUrl;    // 从 application.properties 或 yml 中注入
+    private String baseUrl;
 
     @Value("${dify.datasetsKey}")
-    private String apiKey;     // 从 application.properties 或 yml 中注入
+    private String apiKey;
+
+    @Value("${ragflow.apiKey}")
+    private String ragflowApiKey;
+
+    @Value("${ragflow.endpoint}")
+    private String ragflowBaseUrl;
 
     private CloseableHttpClient httpClient;
 
@@ -228,4 +238,112 @@ public class DifyClient  {
     public void close() throws IOException {
         httpClient.close();
     }
+
+    /*
+     * 上传文档到 Ragflow
+     */
+    public JSONObject ragflowCreateDocumentByFile(List<File> files,String originalFilename, String datasetId) throws IOException {
+        String url = ragflowBaseUrl + "/api/v1/datasets/" + datasetId + "/documents";
+        MultipartEntityBuilder builder = MultipartEntityBuilder.create()
+                .setCharset(StandardCharsets.UTF_8)
+                .setMode(HttpMultipartMode.RFC6532);
+
+        // 添加多个 file
+        for (File file : files) {
+            builder.addBinaryBody("file", file, ContentType.APPLICATION_OCTET_STREAM, originalFilename);
+        }
+        HttpPost post = new HttpPost(url);
+        post.setHeader(AUTH_HEADER, BEARER + ragflowApiKey);
+        post.setEntity(builder.build());
+        return executeJson(post, 200);
+    }
+
+
+    public void deleteRagFlowDocuments(String datasetId, String documentId) throws IOException {
+        String url = String.format("%s/api/v1/datasets/%s/documents", ragflowBaseUrl, datasetId);
+
+        // 用最原始方式构建请求
+        HttpURLConnection conn = (HttpURLConnection) new URL(url).openConnection();
+        conn.setRequestMethod("DELETE");
+        conn.setRequestProperty("Content-Type", "application/json");
+        conn.setRequestProperty("Authorization", "Bearer " + ragflowApiKey);
+        conn.setDoOutput(true);
+
+        // 写入请求体
+        String body = new JSONObject() {{
+            put("ids", Collections.singletonList(documentId));
+        }}.toJSONString();
+
+        try (OutputStream os = conn.getOutputStream()) {
+            os.write(body.getBytes(StandardCharsets.UTF_8));
+        }
+
+        int code = conn.getResponseCode();
+        if (code != 200) {
+            throw new IOException("Delete failed, code: " + code);
+        }
+    }
+
+    public JSONObject ragflowDocumentChunks(String datasetId, List<String> documentIds) throws IOException {
+        String url = String.format("%s/api/v1/datasets/%s/chunks", ragflowBaseUrl, datasetId);
+
+        HttpPost post = new HttpPost(url);
+        post.setHeader("Content-Type", "application/json");
+        post.setHeader(AUTH_HEADER, BEARER + ragflowApiKey);
+
+        JSONObject body = new JSONObject();
+        body.put("document_ids", documentIds);
+
+        post.setEntity(new StringEntity(body.toJSONString(), ContentType.APPLICATION_JSON));
+        return executeJson(post, 200);
+    }
+
+    public void deleteRagflowDocumentChunks(String datasetId, List<String> documentIds) throws IOException {
+        String url = String.format("%s/api/v1/datasets/%s/chunks", ragflowBaseUrl, datasetId);
+
+        HttpURLConnection conn = (HttpURLConnection) new URL(url).openConnection();
+        conn.setRequestMethod("DELETE");
+        conn.setRequestProperty("Content-Type", "application/json");
+        conn.setRequestProperty("Authorization", "Bearer " + ragflowApiKey);
+        conn.setDoOutput(true);
+
+        // 构建 JSON 请求体
+        String body = new JSONObject() {{
+            put("document_ids", documentIds);
+        }}.toJSONString();
+
+        try (OutputStream os = conn.getOutputStream()) {
+            os.write(body.getBytes(StandardCharsets.UTF_8));
+        }
+
+        int responseCode = conn.getResponseCode();
+        if (responseCode != 200) {
+            throw new IOException("DELETE /chunks failed, code: " + responseCode);
+        }
+    }
+
+    public JSONObject getRagflowDocumentStatus(String datasetId, String documentId) throws Exception {
+        String urlString = String.format("%s/api/v1/datasets/%s/documents/?id=%s", ragflowBaseUrl, datasetId, documentId);
+        URL url = new URL(urlString);
+        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
+        conn.setRequestMethod("GET");
+        conn.setRequestProperty("Authorization", "Bearer " + ragflowApiKey);
+        conn.setRequestProperty("Content-Type", "application/json");
+
+        int responseCode = conn.getResponseCode();
+        if (responseCode != HttpURLConnection.HTTP_OK) {
+            throw new RuntimeException("HTTP GET Request Failed with Error code : " + responseCode);
+        }
+
+        BufferedReader in = new BufferedReader(new InputStreamReader(conn.getInputStream()));
+        String inputLine;
+        StringBuilder response = new StringBuilder();
+        while ((inputLine = in.readLine()) != null) {
+            response.append(inputLine);
+        }
+        in.close();
+
+        return JSONObject.parseObject(response.toString());
+    }
+
 }