From a67026f7149c1bdf5007c05b8e8ff759f8ea7bee Mon Sep 17 00:00:00 2001 From: Hz_ Date: Thu, 2 Jul 2026 12:35:10 +0800 Subject: [PATCH] fix(go): agent explore thumbnail loading for multiple doc_ids (#16514) ## Summary - align the Go `/api/v1/thumbnails` endpoint with the frontend request format for repeated `doc_ids` - return thumbnail mappings for multiple documents instead of failing on a single missing document - preserve Python-compatible thumbnail formatting, including base64 thumbnail passthrough --- internal/dao/document.go | 16 ++++++ internal/handler/document.go | 45 +++++++++------- internal/handler/document_test.go | 52 ++++++++++++++++++- internal/service/document.go | 51 +++++++++++++++---- internal/service/document_test.go | 85 +++++++++++++++++++++++++++++++ 5 files changed, 221 insertions(+), 28 deletions(-) diff --git a/internal/dao/document.go b/internal/dao/document.go index 281affd230..af5eeb3491 100644 --- a/internal/dao/document.go +++ b/internal/dao/document.go @@ -226,6 +226,22 @@ func (dao *DocumentDAO) GetByIDs(ids []string) ([]*entity.Document, error) { return documents, nil } +// GetByIDsAndTenantIDs retrieves documents by IDs scoped to knowledgebase owners. +func (dao *DocumentDAO) GetByIDsAndTenantIDs(ids, tenantIDs []string) ([]*entity.Document, error) { + if len(ids) == 0 || len(tenantIDs) == 0 { + return nil, nil + } + var documents []*entity.Document + err := DB.Model(&entity.Document{}). + Joins("JOIN knowledgebase ON document.kb_id = knowledgebase.id"). + Where("document.id IN ? AND knowledgebase.tenant_id IN ? AND knowledgebase.status = ?", ids, tenantIDs, string(entity.StatusValid)). + Find(&documents).Error + if err != nil { + return nil, err + } + return documents, nil +} + // GetByDocumentIDAndDatasetID retrieves a document by document ID and dataset/KB ID. func (dao *DocumentDAO) GetByDocumentIDAndDatasetID(documentID, datasetID string) (*entity.Document, error) { var document entity.Document diff --git a/internal/handler/document.go b/internal/handler/document.go index 14c5389c28..d0d2442917 100644 --- a/internal/handler/document.go +++ b/internal/handler/document.go @@ -52,7 +52,7 @@ type documentServiceIface interface { ListDocuments(page, pageSize int) ([]*service.DocumentResponse, int64, error) ListDocumentsByDatasetID(kbID string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) GetDocumentsByAuthorID(authorID, page, pageSize int) ([]*service.DocumentResponse, int64, error) - GetThumbnail(docID string) (*service.ThumbnailResponse, error) + GetThumbnails(userID string, docIDs []string) (map[string]string, error) GetDocumentImage(imageID string) ([]byte, error) GetMetadataSummary(kbID string, docIDs []string) (map[string]interface{}, error) SetDocumentMetadata(docID string, meta map[string]interface{}) error @@ -168,40 +168,51 @@ func (h *DocumentHandler) GetDocumentByID(c *gin.Context) { // GetThumbnail Get thumbnails for documents. func (h *DocumentHandler) GetThumbnail(c *gin.Context) { - _, errorCode, errorMessage := GetUser(c) + user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { jsonError(c, errorCode, errorMessage) return } - id := c.Query("doc_ids") - if id == "" { - c.JSON(http.StatusBadRequest, gin.H{ - "error": errors.New("invalid document id"), - }) + docIDs := parseThumbnailDocIDs(c) + if len(docIDs) == 0 { + jsonError(c, common.CodeArgumentError, `Lack of "Document ID"`) return } - result, err := h.documentService.GetThumbnail(id) + result, err := h.documentService.GetThumbnails(user.ID, docIDs) if err != nil { - c.JSON(http.StatusNotFound, gin.H{ - "error": fmt.Errorf("thumbnail not found"), - }) + jsonError(c, common.CodeServerError, err.Error()) return } - if result.Thumbnail != nil && *result.Thumbnail != "" { - newThumbURL := fmt.Sprintf("/api/v1/documents/images/%s-%s", result.KbID, *result.Thumbnail) - result.Thumbnail = &newThumbURL - } - c.JSON(http.StatusOK, gin.H{ "code": common.CodeSuccess, - "data": map[string]interface{}{result.ID: result.Thumbnail}, + "data": result, "message": "success", }) } +func parseThumbnailDocIDs(c *gin.Context) []string { + rawValues := c.QueryArray("doc_ids") + seen := make(map[string]struct{}, len(rawValues)) + docIDs := make([]string, 0, len(rawValues)) + + for _, raw := range rawValues { + id := strings.TrimSpace(raw) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + docIDs = append(docIDs, id) + } + + return docIDs +} + // GetDocumentImage returns a document image from object storage. func (h *DocumentHandler) GetDocumentImage(c *gin.Context) { imageID := c.Param("image_id") diff --git a/internal/handler/document_test.go b/internal/handler/document_test.go index b6b8f6852c..22be6fdc82 100644 --- a/internal/handler/document_test.go +++ b/internal/handler/document_test.go @@ -42,6 +42,10 @@ type fakeDocumentService struct { err error stopResult map[string]interface{} stopErr error + thumbnails map[string]string + thumbnailErr error + thumbnailUserID string + thumbnailDocIDs []string metadataSummary map[string]interface{} metadataErr error metadataKBID string @@ -145,8 +149,10 @@ func (f *fakeDocumentService) ListDocumentsByDatasetID(kbID string, page, pageSi func (f *fakeDocumentService) BatchUpdateDocumentStatus(userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { return map[string]interface{}{}, common.CodeSuccess, nil } -func (f *fakeDocumentService) GetThumbnail(docID string) (*service.ThumbnailResponse, error) { - return nil, nil +func (f *fakeDocumentService) GetThumbnails(userID string, docIDs []string) (map[string]string, error) { + f.thumbnailUserID = userID + f.thumbnailDocIDs = append([]string(nil), docIDs...) + return f.thumbnails, f.thumbnailErr } func (f *fakeDocumentService) GetDocumentImage(imageID string) ([]byte, error) { return nil, nil @@ -908,6 +914,48 @@ func TestMetadataSummaryByDataset_Success(t *testing.T) { } } +func TestGetThumbnail_Success(t *testing.T) { + gin.SetMode(gin.TestMode) + fake := &fakeDocumentService{ + thumbnails: map[string]string{ + "doc-1": "/api/v1/documents/images/kb-1-thumb-1.png", + "doc-2": "", + }, + } + h := &DocumentHandler{ + documentService: fake, + } + + c, w := setupGinContextWithUser("GET", "/api/v1/thumbnails?doc_ids=doc-1&doc_ids=doc-2", "") + + h.GetThumbnail(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if len(fake.thumbnailDocIDs) != 2 || fake.thumbnailDocIDs[0] != "doc-1" || fake.thumbnailDocIDs[1] != "doc-2" { + t.Fatalf("unexpected docIDs: %#v", fake.thumbnailDocIDs) + } + if fake.thumbnailUserID != "user-1" { + t.Fatalf("unexpected userID: %s", fake.thumbnailUserID) + } + + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("failed to unmarshal response: %v", err) + } + if resp["code"] != float64(common.CodeSuccess) { + t.Fatalf("expected code %d, got %v", common.CodeSuccess, resp["code"]) + } + data := resp["data"].(map[string]interface{}) + if data["doc-1"] != "/api/v1/documents/images/kb-1-thumb-1.png" { + t.Fatalf("unexpected thumbnail for doc-1: %v", data["doc-1"]) + } + if data["doc-2"] != "" { + t.Fatalf("unexpected thumbnail for doc-2: %v", data["doc-2"]) + } +} + func TestGetDocumentArtifact_Success(t *testing.T) { gin.SetMode(gin.TestMode) h := &DocumentHandler{ diff --git a/internal/service/document.go b/internal/service/document.go index 8165898759..a1328a684a 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -141,6 +141,8 @@ type ThumbnailResponse struct { KbID string `json:"kb_id"` } +const imgBase64Prefix = "data:image/png;base64," + type ArtifactResponse struct { Data []byte ContentType string @@ -874,17 +876,48 @@ func (s *DocumentService) ListDocuments(page, pageSize int) ([]*DocumentResponse return responses, total, nil } -func (s *DocumentService) GetThumbnail(docID string) (*ThumbnailResponse, error) { - document, err := s.documentDAO.GetByID(docID) - if err != nil { - return nil, err +func (s *DocumentService) GetThumbnails(userID string, docIDs []string) (map[string]string, error) { + if len(docIDs) == 0 { + return map[string]string{}, nil } - var result ThumbnailResponse - result.ID = document.ID - result.Thumbnail = document.Thumbnail - result.KbID = document.KbID - return &result, nil + tenantIDs := []string{userID} + if userID != "" { + ids, err := dao.NewUserTenantDAO().GetTenantIDsByUserID(userID) + if err != nil { + return nil, fmt.Errorf("failed to fetch user tenants: %w", err) + } + tenantIDs = append(tenantIDs, ids...) + } + + documents, err := s.documentDAO.GetByIDsAndTenantIDs(docIDs, tenantIDs) + if err != nil { + return nil, fmt.Errorf("failed to fetch document thumbnails: %w", err) + } + + result := make(map[string]string, len(documents)) + for _, document := range documents { + if document == nil { + continue + } + + thumbnail := "" + if document.Thumbnail != nil && *document.Thumbnail != "" { + if strings.HasPrefix(*document.Thumbnail, imgBase64Prefix) { + thumbnail = *document.Thumbnail + } else { + thumbnail = fmt.Sprintf( + "/api/v1/documents/images/%s-%s", + document.KbID, + *document.Thumbnail, + ) + } + } + + result[document.ID] = thumbnail + } + + return result, nil } func (s *DocumentService) BatchUpdateDocumentStatus(userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) { diff --git a/internal/service/document_test.go b/internal/service/document_test.go index 1237b5f875..3c5889d652 100644 --- a/internal/service/document_test.go +++ b/internal/service/document_test.go @@ -2230,3 +2230,88 @@ func TestGetDocumentArtifact_AuthGate(t *testing.T) { t.Errorf("user-2 with unrelated session: want ErrArtifactNotFound, got %v", err) } } + +func TestGetThumbnails_AlignsWithPythonFormatting(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + + if err := db.AutoMigrate(&entity.Document{}, &entity.Knowledgebase{}, &entity.UserTenant{}); err != nil { + t.Fatalf("migrate: %v", err) + } + + insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) + insertTestKB(t, "kb-2", "tenant-1", 0, 0, 0) + insertTestKB(t, "kb-other", "tenant-other", 0, 0, 0) + if err := db.Create(&entity.UserTenant{ + ID: "user-1_tenant-1", + UserID: "user-1", + TenantID: "tenant-1", + Role: "owner", + InvitedBy: "user-1", + Status: sptr("1"), + }).Error; err != nil { + t.Fatalf("seed user tenant: %v", err) + } + + base64Thumb := "data:image/png;base64,AAAA" + fileThumb := "thumb.png" + otherThumb := "secret.png" + if err := db.Create(&entity.Document{ + ID: "doc-file", + KbID: "kb-1", + Thumbnail: &fileThumb, + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + SourceType: "local", + Type: "pdf", + CreatedBy: "user-1", + Suffix: "png", + }).Error; err != nil { + t.Fatalf("seed file thumbnail doc: %v", err) + } + if err := db.Create(&entity.Document{ + ID: "doc-base64", + KbID: "kb-2", + Thumbnail: &base64Thumb, + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + SourceType: "local", + Type: "pdf", + CreatedBy: "user-1", + Suffix: "png", + }).Error; err != nil { + t.Fatalf("seed base64 thumbnail doc: %v", err) + } + if err := db.Create(&entity.Document{ + ID: "doc-other", + KbID: "kb-other", + Thumbnail: &otherThumb, + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + SourceType: "local", + Type: "pdf", + CreatedBy: "user-other", + Suffix: "png", + }).Error; err != nil { + t.Fatalf("seed other tenant thumbnail doc: %v", err) + } + + svc := testDocumentService(t) + got, err := svc.GetThumbnails("user-1", []string{"doc-file", "doc-base64", "doc-other", "missing-doc"}) + if err != nil { + t.Fatalf("GetThumbnails failed: %v", err) + } + + if got["doc-file"] != "/api/v1/documents/images/kb-1-thumb.png" { + t.Fatalf("unexpected file thumbnail: %q", got["doc-file"]) + } + if got["doc-base64"] != base64Thumb { + t.Fatalf("unexpected base64 thumbnail: %q", got["doc-base64"]) + } + if _, ok := got["missing-doc"]; ok { + t.Fatalf("did not expect missing doc in result: %#v", got) + } + if _, ok := got["doc-other"]; ok { + t.Fatalf("did not expect other tenant doc in result: %#v", got) + } +}