package ragflow import ( "bytes" "context" "encoding/json" "fmt" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/net/gclient" "github.com/gogf/gf/v2/os/gctx" "github.com/gogf/gf/v2/os/glog" "github.com/gogf/gf/v2/util/gconv" "io" "mime/multipart" "net/http" "regexp" "sync" ) var ( once sync.Once internalClient *Client // 内部私有客户端 ) // 32位小写十六进制正则校验 var pipelineIDRegex = regexp.MustCompile(`^[0-9a-f]{32}$`) // Client 包装 gclient.Client type Client struct { cli *gclient.Client // 原生GoFrame HTTP客户端 address string // RAGFlow服务地址 apiKey string // 认证密钥 } // Init 初始化客户端 func Init(ctx context.Context, addr, apiKey string) *Client { once.Do(func() { internalClient = newClient(ctx, addr, apiKey) }) return internalClient } // GetHttpClient 获取原生 gclient.Client func GetHttpClient() *Client { return internalClient } // newClient 创建客户端实例 func newClient(ctx context.Context, addr, apiKey string) *Client { // 初始化GoFrame HTTP客户端 cli := g.Client() // 设置全局默认请求头 cli.SetHeader("Content-Type", "application/json") cli.SetHeader("Authorization", "Bearer "+apiKey) // 初始化校验 if addr == "" || apiKey == "" { panic("ragflow client init failed: address or apiKey is empty") } return &Client{ cli: cli, address: addr, apiKey: apiKey, } } // -------------------------- 工具函数 -------------------------- func Log(ctx context.Context, v ...interface{}) { glog.Info(ctx, v...) } func NewCtx() context.Context { return gctx.New() } // -------------------------- 公共结构体 -------------------------- type CommonResp struct { Code int `json:"code"` Msg string `json:"message,omitempty"` } // -------------------------- 1. 数据集接口(原有) -------------------------- type CreateDatasetReq struct { Name string `json:"name"` ChunkMethod string `json:"chunk_method,omitempty"` Description string `json:"description,omitempty"` EmbeddingModel string `json:"embedding_model,omitempty"` //Avatar string `json:"avatar,omitempty"` //Permission string `json:"permission,omitempty"` //ParserConfig interface{} `json:"parser_config,omitempty"` //ParseType int `json:"parse_type,omitempty"` //PipelineID string `json:"pipeline_id,omitempty"` } type CreateDatasetResp struct { Code int `json:"code"` Data DatasetData `json:"data"` } type DatasetData struct { ID string `json:"id"` Name string `json:"name"` ChunkMethod string `json:"chunk_method"` EmbeddingModel string `json:"embedding_model"` Permission string `json:"permission"` ParserConfig interface{} `json:"parser_config"` CreateTime int64 `json:"create_time"` DocumentCount int `json:"document_count"` ChunkCount int `json:"chunk_count"` } type DeleteDatasetReq struct { IDs []string `json:"ids"` } type UpdateDatasetReq struct { Name string `json:"name,omitempty"` Language string `json:"language"` Avatar string `json:"avatar,omitempty"` Description string `json:"description,omitempty"` EmbeddingModel string `json:"embedding_model,omitempty"` Permission string `json:"permission,omitempty"` ChunkMethod string `json:"chunk_method,omitempty"` ParserConfig interface{} `json:"parser_config,omitempty"` Pagerank int `json:"pagerank,omitempty"` } // 创建知识库 func (r *Client) CreateDataset(ctx context.Context, req *CreateDatasetReq) (*CreateDatasetResp, error) { if req.Name == "" { return nil, fmt.Errorf("name 是必填参数") } url := fmt.Sprintf("%s/api/v1/datasets", r.address) var resp CreateDatasetResp err := r.cli.PostVar(ctx, url, req).Scan(&resp) if err != nil { return nil, fmt.Errorf("请求失败: %v", err) } if resp.Code != 0 { return nil, fmt.Errorf("接口错误 code=%d", resp.Code) } return &resp, nil } // DeleteDataset 删除知识库 func (r *Client) DeleteDataset(ctx context.Context, ids []string) (*CommonResp, error) { url := fmt.Sprintf("%s/api/v1/datasets", r.address) req := DeleteDatasetReq{IDs: ids} var resp CommonResp err := r.cli.DeleteVar(ctx, url, req).Scan(&resp) if err != nil { return nil, fmt.Errorf("请求失败: %v", err) } return &resp, nil } // UpdateDataset 更新知识库 func (r *Client) UpdateDataset(ctx context.Context, datasetID string, req *UpdateDatasetReq) (*CommonResp, error) { if datasetID == "" { return nil, fmt.Errorf("datasetID 不能为空") } url := fmt.Sprintf("%s/api/v1/datasets/%s", r.address, datasetID) var resp CommonResp err := r.cli.PutVar(ctx, url, req).Scan(&resp) if err != nil { return nil, fmt.Errorf("请求失败: %v", err) } return &resp, nil } // -------------------------- 2. 文档接口(原有) -------------------------- type UploadFileReq struct { File []*FileInfo `json:"file"` } type FileInfo struct { FileName string `json:"file_name"` // 文件名称 Url string `json:"url"` // 文件地址 } type UploadDocumentResp struct { Code int `json:"code"` Data []DocumentDetail `json:"data"` Message string `json:"message"` // 失败时返回 } type DocumentDetail struct { ID string `json:"id"` Name string `json:"name"` DatasetID string `json:"dataset_id"` Size int64 `json:"size"` Run string `json:"run"` ChunkMethod string `json:"chunk_method"` ParserConfig interface{} `json:"parser_config"` Status string `json:"status"` } type UpdateDocumentReq struct { Name string `json:"name,omitempty"` MetaFields interface{} `json:"meta_fields,omitempty"` ChunkMethod string `json:"chunk_method,omitempty"` ParserConfig interface{} `json:"parser_config,omitempty"` Enabled int `json:"enabled,omitempty"` } type UpdateDocumentResp struct { Code int `json:"code"` Message string `json:"message"` // 失败时返回 Data DocumentDetail `json:"data"` } type ListDocumentReq struct { Page int `json:"page,omitempty"` PageSize int `json:"page_size,omitempty"` OrderBy string `json:"orderby,omitempty"` Desc bool `json:"desc,omitempty"` Keywords string `json:"keywords,omitempty"` DocumentID string `json:"id,omitempty"` DocumentName string `json:"name,omitempty"` CreateTimeFrom int64 `json:"create_time_from,omitempty"` CreateTimeTo int64 `json:"create_time_to,omitempty"` Suffix string `json:"suffix,omitempty"` Run string `json:"run,omitempty"` MetadataCondition string `json:"metadata_condition,omitempty"` } type ListDocumentResp struct { Code int `json:"code"` Data struct { Docs []DocumentDetail `json:"docs"` Total int `json:"total_datasets"` } `json:"data"` } type DeleteDocumentsReq struct { IDs []string `json:"ids"` } type ParseDocumentReq struct { DocumentIDs []string `json:"document_ids"` } //func (r *Client) UploadDocument(ctx context.Context, datasetID, fileName string, url string) error { // // 参数校验 // remoteFileUrl := url // apiUrl := fmt.Sprintf("%s/api/v1/datasets/%s/documents", r.address, datasetID) // fieldName := "file" // // resp, err := http.Get(remoteFileUrl) // if err != nil { // panic("下载远程文件失败:" + err.Error()) // } // defer resp.Body.Close() // // body := &bytes.Buffer{} // writer := multipart.NewWriter(body) // // formFile, _ := writer.CreateFormFile(fieldName, "xinfengjinghua.pdf") // io.Copy(formFile, resp.Body) // _ = writer.Close() // // req, _ := http.NewRequest("POST", apiUrl, body) // req.Header.Set("Content-Type", writer.FormDataContentType()) // req.Header.Set("Authorization", "Bearer "+"ragflow-nkGJCVuFSxgt7X6AkzuOXjx_9Q3hN3zAUP9LNpmJOFk") // // client := &http.Client{} // uploadResp, err := client.Do(req) // if err != nil { // panic("上传失败:" + err.Error()) // } // defer uploadResp.Body.Close() // // result, _ := io.ReadAll(uploadResp.Body) // println("上传状态:", uploadResp.Status) // println("接口返回:", string(result)) // // return nil //} func (r *Client) UploadDocument(ctx context.Context, datasetID string, req UploadFileReq) (*UploadDocumentResp, error) { // 参数校验 if datasetID == "" { return nil, fmt.Errorf("datasetID 不能为空") } body := &bytes.Buffer{} writer := multipart.NewWriter(body) for _, v := range req.File { //下载远程文件 f, err := r.cli.Get(ctx, v.Url) if err != nil { return nil, fmt.Errorf("下载文件失败: %v", err) } part, err := writer.CreateFormFile("file", v.FileName) if err != nil { return nil, fmt.Errorf("创建表单失败: %v", err) } _, err = io.Copy(part, f.Response.Body) if err != nil { return nil, fmt.Errorf("写入文件流失败: %v", err) } _ = f.Response.Body.Close() } err := writer.Close() if err != nil { return nil, fmt.Errorf("关闭表单失败: %v", err) } r.cli.SetHeader("Content-Type", writer.FormDataContentType()) // 拼接上传接口地址 apiUrl := fmt.Sprintf("%s/api/v1/datasets/%s/documents", r.address, datasetID) var resp UploadDocumentResp if err = r.cli.PostVar(ctx, apiUrl, body).Scan(&resp); err != nil { return nil, fmt.Errorf("解析响应失败: %v", err) } if resp.Code != 0 { return nil, fmt.Errorf("上传失败[%d]: %s", resp.Code, resp.Message) } r.cli.SetHeader("Content-Type", "application/json") return &resp, nil } // UpdateDocument 更新文档 func (r *Client) UpdateDocument(ctx context.Context, datasetID, documentID string, req *UpdateDocumentReq) (*UpdateDocumentResp, error) { url := fmt.Sprintf("%s/api/v1/datasets/%s/documents/%s", r.address, datasetID, documentID) var resp UpdateDocumentResp err := r.cli.PutVar(ctx, url, req).Scan(&resp) if err != nil { return nil, fmt.Errorf("更新失败: %v", err) } return &resp, nil } // ListDocuments 列出文档 func (r *Client) ListDocuments(ctx context.Context, datasetID string, req *ListDocumentReq) (*ListDocumentResp, error) { url := fmt.Sprintf("%s/api/v1/datasets/%s/documents", r.address, datasetID) params := gconv.Map(req) var resp ListDocumentResp err := r.cli.GetVar(ctx, url, params).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // DeleteDocuments 删除文档 func (r *Client) DeleteDocuments(ctx context.Context, datasetID string, ids []string) (*CommonResp, error) { url := fmt.Sprintf("%s/api/v1/datasets/%s/documents", r.address, datasetID) req := DeleteDocumentsReq{IDs: ids} var resp CommonResp err := r.cli.DeleteVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // ParseDocuments 解析文档 func (r *Client) ParseDocuments(ctx context.Context, datasetID string, ids []string) (*CommonResp, error) { if len(ids) == 0 { return nil, fmt.Errorf("document_ids 不能为空") } url := fmt.Sprintf("%s/api/v1/datasets/%s/chunks", r.address, datasetID) req := ParseDocumentReq{DocumentIDs: ids} var resp CommonResp err := r.cli.PostVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // -------------------------- 3. 对话助手接口(新增) -------------------------- // LLM配置 type ChatLLM struct { ModelName string `json:"model_name,omitempty"` ModelType string `json:"model_type,omitempty"` // chat/image2text Temperature float64 `json:"temperature,omitempty"` TopP float64 `json:"top_p,omitempty"` PresencePenalty float64 `json:"presence_penalty,omitempty"` FrequencyPenalty float64 `json:"frequency_penalty,omitempty"` } // Prompt配置 type ChatPrompt struct { SimilarityThreshold float64 `json:"similarity_threshold,omitempty"` KeywordsSimilarityWeight float64 `json:"keywords_similarity_weight,omitempty"` TopN int `json:"top_n,omitempty"` Variables []interface{} `json:"variables,omitempty"` RerankModel string `json:"rerank_model,omitempty"` EmptyResponse string `json:"empty_response,omitempty"` Opener string `json:"opener,omitempty"` ShowQuote bool `json:"show_quote,omitempty"` Prompt string `json:"prompt,omitempty"` } // 创建对话助手 type CreateChatReq struct { Name string `json:"name"` // 必填 Avatar string `json:"avatar,omitempty"` DatasetIDs []string `json:"dataset_ids,omitempty"` LLM *ChatLLM `json:"llm,omitempty"` Prompt *ChatPrompt `json:"prompt,omitempty"` } type CreateChatResp struct { Code int `json:"code"` Data ChatInfo `json:"data"` } type ChatInfo struct { ID string `json:"id"` Name string `json:"name"` } // 更新对话助手 type UpdateChatReq struct { Name string `json:"name,omitempty"` Avatar string `json:"avatar,omitempty"` DatasetIDs []string `json:"dataset_ids,omitempty"` LLM *ChatLLM `json:"llm,omitempty"` Prompt *ChatPrompt `json:"prompt,omitempty"` } // 列表助手 type ListChatReq struct { Page int `json:"page,omitempty"` PageSize int `json:"page_size,omitempty"` OrderBy string `json:"orderby,omitempty"` Desc bool `json:"desc,omitempty"` ID string `json:"id,omitempty"` Name string `json:"name,omitempty"` } type ListChatResp struct { Code int `json:"code"` Data []interface{} `json:"data"` } // 删除助手 type DeleteChatReq struct { IDs []string `json:"ids"` } // CreateChat 创建对话助手 func (r *Client) CreateChat(ctx context.Context, req *CreateChatReq) (*CreateChatResp, error) { if req.Name == "" { return nil, fmt.Errorf("name 是必填参数") } url := fmt.Sprintf("%s/api/v1/chats", r.address) r.cli.SetHeader("Content-Type", "application/json") var resp CreateChatResp err := r.cli.PostVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // UpdateChat 更新对话助手 func (r *Client) UpdateChat(ctx context.Context, chatID string, req *UpdateChatReq) (*CommonResp, error) { if chatID == "" { return nil, fmt.Errorf("chatID 不能为空") } url := fmt.Sprintf("%s/api/v1/chats/%s", r.address, chatID) r.cli.SetHeader("Content-Type", "application/json") var resp CommonResp err := r.cli.PutVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // ListChats 列出助手 func (r *Client) ListChats(ctx context.Context, req *ListChatReq) (*ListChatResp, error) { url := fmt.Sprintf("%s/api/v1/chats", r.address) params := gconv.Map(req) var resp ListChatResp err := r.cli.GetVar(ctx, url, params).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // DeleteChats 删除助手 func (r *Client) DeleteChats(ctx context.Context, ids []string) (*CommonResp, error) { url := fmt.Sprintf("%s/api/v1/chats", r.address) r.cli.SetHeader("Content-Type", "application/json") req := DeleteChatReq{IDs: ids} var resp CommonResp err := r.cli.DeleteVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // -------------------------- 4. 会话管理接口(新增) -------------------------- type CreateSessionReq struct { Name string `json:"name"` // 必填 UserID string `json:"user_id,omitempty"` } type CreateSessionResp struct { Code int `json:"code"` Data SessionInfo `json:"data"` } type SessionInfo struct { Id string `json:"id"` Message struct { Content string `json:"content"` Role string `json:"role"` } `json:"message"` Name string `json:"name"` } type UpdateSessionReq struct { Name string `json:"name,omitempty"` UserID string `json:"user_id,omitempty"` } type ListSessionReq struct { Page int `json:"page,omitempty"` PageSize int `json:"page_size,omitempty"` OrderBy string `json:"orderby,omitempty"` Desc bool `json:"desc,omitempty"` ID string `json:"id,omitempty"` Name string `json:"name,omitempty"` UserID string `json:"user_id,omitempty"` } type DeleteSessionReq struct { IDs []string `json:"ids"` } // CreateSession 创建会话 func (r *Client) CreateSession(ctx context.Context, chatID string, req *CreateSessionReq) (*CreateSessionResp, error) { if chatID == "" || req.Name == "" { return nil, fmt.Errorf("参数不能为空") } url := fmt.Sprintf("%s/api/v1/chats/%s/sessions", r.address, chatID) var resp CreateSessionResp err := r.cli.PostVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // UpdateSession 更新会话 func (r *Client) UpdateSession(ctx context.Context, chatID, sessionID string, req *UpdateSessionReq) (*CommonResp, error) { url := fmt.Sprintf("%s/api/v1/chats/%s/sessions/%s", r.address, chatID, sessionID) var resp CommonResp err := r.cli.PutVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // ListSessions 列出会话 func (r *Client) ListSessions(ctx context.Context, chatID string, req *ListSessionReq) (*CommonResp, error) { url := fmt.Sprintf("%s/api/v1/chats/%s/sessions", r.address, chatID) params := gconv.Map(req) var resp CommonResp err := r.cli.GetVar(ctx, url, params).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // DeleteSessions 删除会话 func (r *Client) DeleteSessions(ctx context.Context, chatID string, ids []string) (*CommonResp, error) { url := fmt.Sprintf("%s/api/v1/chats/%s/sessions", r.address, chatID) req := DeleteSessionReq{IDs: ids} var resp CommonResp err := r.cli.DeleteVar(ctx, url, req).Scan(&resp) if err != nil { return nil, err } return &resp, nil } // -------------------------- 5. AI对话接口(流式/非流式 新增) -------------------------- // 元数据过滤条件 type MetaCondition struct { Logic string `json:"logic,omitempty"` // and/or Conditions []Condition `json:"conditions,omitempty"` } type Condition struct { Name string `json:"name"` ComparisonOperator string `json:"comparison_operator"` Value interface{} `json:"value,omitempty"` } // 对话请求 type ChatCompletionReq struct { Question string `json:"question"` // 必填 Stream bool `json:"stream"` SessionID string `json:"session_id,omitempty"` UserID string `json:"user_id,omitempty"` MetadataCondition *MetaCondition `json:"metadata_condition,omitempty"` } type Completion struct { Answer string `json:"answer"` Reference interface{} `json:"reference"` AudioBinary interface{} `json:"audio_binary"` Id string `json:"id"` SessionId string `json:"session_id"` } // 对话请求 type ChatCompletionResp struct { Code int `json:"code"` Message string `json:"message"` Data Completion `json:"data"` } // ChatCompletions 对话 func (r *Client) ChatCompletions(ctx context.Context, chatID string, req *ChatCompletionReq) (*ChatCompletionResp, error) { if chatID == "" || req.Question == "" { return nil, fmt.Errorf("chatID 和 question 不能为空") } url := fmt.Sprintf("%s/api/v1/chats/%s/completions", r.address, chatID) req.Stream = false var resp ChatCompletionResp d, err := r.cli.Post(ctx, url, req) if err != nil { return nil, err } err = json.Unmarshal(d.ReadAll(), &resp) return &resp, nil } func (r *Client) ChatCompletionsStream(ctx context.Context, chatID string, req *ChatCompletionReq) (io.ReadCloser, error) { if chatID == "" || req.Question == "" { return nil, fmt.Errorf("chatID 和 question 不能为空") } // 强制开启流式 req.Stream = true url := fmt.Sprintf("%s/api/v1/chats/%s/completions", r.address, chatID) r.cli.Timeout(0) r.cli.Transport = &http.Transport{ DisableCompression: true, } resp, err := r.cli.Post(ctx, url, req) if err != nil { return nil, err } return resp.Body, nil }