mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-19 22:21:04 +08:00
feat(go-api): migrate agent file download handler to Go with strict P… (#15769)
## What does this PR do? This PR migrates the Agent Temporary File Download endpoint (`GET /api/v1/agents/download`) from the Python backend to the Go backend, optimizing the data retrieval flow and maintaining strict functional parity. It also fixes a persistent parsing error in the Sandbox code execution node. ## Checklist - [x] Code logic matches Python implementation - [x] All local unit tests passed - [x] No breaking changes to existing router interfaces Co-authored-by: Yingfeng <yingfeng.zhang@gmail.com>
This commit is contained in:
@@ -22,6 +22,8 @@ import (
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"ragflow/internal/utility"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
@@ -36,19 +38,24 @@ import (
|
||||
|
||||
// AgentHandler agent handler
|
||||
// fileUploader is the subset of FileService used by agent handlers.
|
||||
type fileUploader interface {
|
||||
type agentFileService interface {
|
||||
UploadFile(tenantID, parentID string, files []*multipart.FileHeader) ([]map[string]interface{}, error)
|
||||
DownloadAgentFile(tenantID, location string) ([]byte, error)
|
||||
}
|
||||
|
||||
// AgentHandler agent handler
|
||||
type AgentHandler struct {
|
||||
agentService *service.AgentService
|
||||
fileService fileUploader
|
||||
fileService agentFileService
|
||||
}
|
||||
|
||||
// NewAgentHandler create agent handler
|
||||
|
||||
func NewAgentHandler(agentService *service.AgentService, fileService *service.FileService) *AgentHandler {
|
||||
return &AgentHandler{agentService: agentService, fileService: fileService}
|
||||
return &AgentHandler{
|
||||
agentService: agentService,
|
||||
fileService: fileService,
|
||||
}
|
||||
}
|
||||
|
||||
// ListAgents lists agent canvases for the current user.
|
||||
@@ -655,14 +662,42 @@ func (h *AgentHandler) UpdateAgentTags(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// GetPrompts returns the default prompts used by the agent.
|
||||
// @Summary Get Agent Prompts
|
||||
// @Description Returns the default prompts used by the agent, such as task analysis, plan generation, reflection, and citation guidelines.
|
||||
// @Tags agents
|
||||
// @Produce json
|
||||
// @Security ApiKeyAuth
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Router /api/v1/agents/prompts [get]
|
||||
func (h *AgentHandler) DownloadAgentFile(c *gin.Context) {
|
||||
user, errorCode, errorMessage := GetUser(c)
|
||||
if errorCode != common.CodeSuccess {
|
||||
jsonError(c, errorCode, errorMessage)
|
||||
return
|
||||
}
|
||||
|
||||
fileID := c.Query("id")
|
||||
if fileID == "" {
|
||||
jsonError(c, common.CodeArgumentError, "id parameter is required")
|
||||
return
|
||||
}
|
||||
|
||||
blob, err := h.fileService.DownloadAgentFile(user.ID, fileID)
|
||||
if err != nil {
|
||||
jsonError(c, common.CodeServerError, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
ext := utility.GetFileExtension(fileID)
|
||||
contentType := utility.GetContentType(ext, "")
|
||||
if contentType != "" {
|
||||
c.Header("Content-Type", contentType)
|
||||
} else {
|
||||
c.Header("Content-Type", "application/octet-stream")
|
||||
}
|
||||
|
||||
if utility.ShouldForceAttachment(ext, contentType) {
|
||||
c.Header("X-Content-Type-Options", "nosniff")
|
||||
encodedName := url.QueryEscape(fileID)
|
||||
c.Header("Content-Disposition", "attachment; filename*=UTF-8''"+encodedName)
|
||||
}
|
||||
|
||||
c.Data(http.StatusOK, contentType, blob)
|
||||
}
|
||||
|
||||
func (h *AgentHandler) GetPrompts(c *gin.Context) {
|
||||
if _, errorCode, errorMessage := GetUser(c); errorCode != common.CodeSuccess {
|
||||
jsonError(c, errorCode, errorMessage)
|
||||
|
||||
@@ -18,6 +18,7 @@ package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -979,9 +980,71 @@ func TestListAgentTemplates_RequiresAuth(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
type fakeAgentFileService struct {
|
||||
blob []byte
|
||||
err error
|
||||
}
|
||||
|
||||
func (f *fakeAgentFileService) UploadFile(tenantID, parentID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (f *fakeAgentFileService) DownloadAgentFile(tenantID, location string) ([]byte, error) {
|
||||
return f.blob, f.err
|
||||
}
|
||||
|
||||
func TestDownloadAgentFile_Success(t *testing.T) {
|
||||
c, w, _ := setupGinContextWithUserAndDB(t, http.MethodGet, "/api/v1/agents/download?id=test-file.pdf")
|
||||
|
||||
fakeFileSvc := &fakeAgentFileService{
|
||||
blob: []byte("test content"),
|
||||
err: nil,
|
||||
}
|
||||
|
||||
h := &AgentHandler{
|
||||
agentService: service.NewAgentService(),
|
||||
fileService: fakeFileSvc,
|
||||
}
|
||||
|
||||
h.DownloadAgentFile(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
if w.Header().Get("Content-Type") != "application/pdf" {
|
||||
t.Errorf("expected Content-Type application/pdf, got %s", w.Header().Get("Content-Type"))
|
||||
}
|
||||
|
||||
if w.Body.String() != "test content" {
|
||||
t.Errorf("expected 'test content', got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadAgentFile_MissingID(t *testing.T) {
|
||||
c, w, _ := setupGinContextWithUserAndDB(t, http.MethodGet, "/api/v1/agents/download")
|
||||
|
||||
h := &AgentHandler{
|
||||
agentService: service.NewAgentService(),
|
||||
fileService: &fakeAgentFileService{},
|
||||
}
|
||||
|
||||
h.DownloadAgentFile(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200 (json error return), got %d", w.Code)
|
||||
}
|
||||
|
||||
var resp map[string]interface{}
|
||||
json.Unmarshal(w.Body.Bytes(), &resp)
|
||||
if code, _ := resp["code"].(float64); code != float64(common.CodeArgumentError) {
|
||||
t.Errorf("expected code 102, got %v", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetPrompts_Success(t *testing.T) {
|
||||
c, w, _ := setupGinContextWithUserAndDB(t, http.MethodGet, "/api/v1/agents/prompts")
|
||||
|
||||
|
||||
// Create handler with fake or real service.
|
||||
h := NewAgentHandler(service.NewAgentService(), nil)
|
||||
h.GetPrompts(c)
|
||||
|
||||
@@ -235,4 +235,7 @@ func TestUploadAgentFileHandler_TeamMemberTenant(t *testing.T) {
|
||||
}
|
||||
|
||||
// sp returns a pointer to the given string.
|
||||
func sp(s string) *string { return &s }
|
||||
func sp(s string) *string { return &s }
|
||||
func (f *fakeUploadFileService) DownloadAgentFile(tenantID, location string) ([]byte, error) {
|
||||
panic("DownloadAgentFile should not be called during upload tests")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user