Initial commit: GovAI 政务AI平台

server/pkg/dify/knowledge.go

@@ -0,0 +1,96 @@
+package dify
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/http"
+)
+
+// CreateDataset creates a new knowledge base (dataset) in Dify.
+func (c *Client) CreateDataset(ctx context.Context, apiKey string, req *DatasetCreateRequest) (*Dataset, error) {
+	var result Dataset
+	if err := c.doJSON(ctx, "POST", "/datasets", apiKey, req, &result); err != nil {
+		return nil, err
+	}
+	return &result, nil
+}
+
+// DeleteDataset deletes a knowledge base.
+func (c *Client) DeleteDataset(ctx context.Context, apiKey, datasetID string) error {
+	return c.doJSON(ctx, "DELETE", "/datasets/"+datasetID, apiKey, nil, nil)
+}
+
+// UploadDocument uploads a file to a dataset for indexing.
+func (c *Client) UploadDocument(ctx context.Context, apiKey, datasetID string, filename string, fileReader io.Reader) (*DocumentIndexingStatus, error) {
+	var buf bytes.Buffer
+	writer := multipart.NewWriter(&buf)
+
+	part, err := writer.CreateFormFile("file", filename)
+	if err != nil {
+		return nil, fmt.Errorf("create form file: %w", err)
+	}
+	if _, err := io.Copy(part, fileReader); err != nil {
+		return nil, fmt.Errorf("copy file: %w", err)
+	}
+
+	// indexing mode
+	if err := writer.WriteField("indexing_technique", "high_quality"); err != nil {
+		return nil, err
+	}
+	if err := writer.WriteField("process_rule", `{"mode": "automatic"}`); err != nil {
+		return nil, err
+	}
+	writer.Close()
+
+	path := fmt.Sprintf("/datasets/%s/document/create_by_file", datasetID)
+	req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+path, &buf)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Authorization", "Bearer "+apiKey)
+	req.Header.Set("Content-Type", writer.FormDataContentType())
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode >= 400 {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("upload failed (status %d): %s", resp.StatusCode, string(body))
+	}
+
+	// Dify returns document info in response
+	var result struct {
+		Document DocumentIndexingStatus `json:"document"`
+	}
+	if err := decodeJSON(resp.Body, &result); err != nil {
+		return nil, err
+	}
+	return &result.Document, nil
+}
+
+// DeleteDocument deletes a document from a dataset.
+func (c *Client) DeleteDocument(ctx context.Context, apiKey, datasetID, documentID string) error {
+	path := fmt.Sprintf("/datasets/%s/documents/%s", datasetID, documentID)
+	return c.doJSON(ctx, "DELETE", path, apiKey, nil, nil)
+}
+
+// GetDocumentIndexingStatus checks the indexing status of a document.
+func (c *Client) GetDocumentIndexingStatus(ctx context.Context, apiKey, datasetID, batch string) (*DocumentIndexingStatus, error) {
+	path := fmt.Sprintf("/datasets/%s/documents/%s/indexing-status", datasetID, batch)
+	var result DocumentIndexingStatus
+	if err := c.doJSON(ctx, "GET", path, apiKey, nil, &result); err != nil {
+		return nil, err
+	}
+	return &result, nil
+}
+
+func decodeJSON(r io.Reader, v any) error {
+	return json.NewDecoder(r).Decode(v)
+}