freedakgmail
97 lines
2.9 KiB
Go
97 lines
2.9 KiB
Go
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)
|
|
}
|