From 4a9977b11fbf2f4fa7072abedbaa777c8cb64f49 Mon Sep 17 00:00:00 2001 From: wizardchen Date: Sat, 28 Feb 2026 18:11:37 +0800 Subject: [PATCH] feat: enhance document loading and state management in doc-content component - Introduced a loading state (`loadingChunks`) to manage document chunk loading more effectively. - Updated the scroll handling logic to prevent multiple requests while loading chunks. - Adjusted the pagination logic to ensure correct page increments based on the total document count. This update improves the user experience by ensuring smoother document loading and better state management during scrolling interactions. --- frontend/src/components/doc-content.vue | 20 ++-- internal/application/service/agent_service.go | 4 +- .../service/chat_pipline/memory.go | 14 +-- .../application/service/chat_pipline/merge.go | 4 +- .../service/chat_pipline/search.go | 2 +- .../service/chat_pipline/search_entity.go | 4 +- internal/application/service/chunk.go | 18 +-- internal/application/service/custom_agent.go | 12 +- internal/application/service/evaluation.go | 4 +- internal/application/service/file/minio.go | 85 ++++++------- .../application/service/image_multimodal.go | 17 +-- internal/application/service/knowledgebase.go | 20 ++-- internal/application/service/message.go | 10 +- internal/application/service/model.go | 12 +- internal/application/service/session.go | 26 ++-- internal/application/service/tag.go | 8 +- internal/handler/initialization.go | 2 +- internal/handler/knowledge.go | 68 +---------- internal/handler/organization.go | 2 +- internal/handler/session/helpers.go | 4 +- internal/handler/tag.go | 2 +- internal/handler/tenant.go | 20 ++-- internal/infrastructure/docparser/helpers.go | 110 +++++++++++++++++ .../docparser/image_resolver.go | 20 +--- .../docparser/mineru_cloud_converter.go | 85 +------------ .../docparser/mineru_converter.go | 113 ++++-------------- internal/types/context_helpers.go | 47 ++++++++ internal/utils/httputil.go | 32 +++++ 28 files changed, 351 insertions(+), 414 deletions(-) create mode 100644 internal/infrastructure/docparser/helpers.go create mode 100644 internal/types/context_helpers.go create mode 100644 internal/utils/httputil.go diff --git a/frontend/src/components/doc-content.vue b/frontend/src/components/doc-content.vue index de189e97..ca4a0a71 100644 --- a/frontend/src/components/doc-content.vue +++ b/frontend/src/components/doc-content.vue @@ -54,6 +54,7 @@ marked.use({ }); const renderer = new marked.Renderer(); let page = 1; +let loadingChunks = false; let doc = null; let down = ref() let mdContentWrap = ref() @@ -136,9 +137,13 @@ onMounted(() => { doc.addEventListener('scroll', handleDetailsScroll); }) }) -onUpdated(() => { - page = 1 -}) +watch(() => props.details?.id, () => { + page = 1; + loadingChunks = false; +}); +watch(() => props.details?.md?.length, () => { + loadingChunks = false; +}); onUnmounted(() => { doc.removeEventListener('scroll', handleDetailsScroll); }) @@ -559,12 +564,13 @@ const downloadFile = () => { }); }; const handleDetailsScroll = () => { - if (doc) { - let pageNum = Math.ceil(props.details.total / 20); + if (doc && !loadingChunks) { + let pageNum = Math.ceil(props.details.total / 25); const { scrollTop, scrollHeight, clientHeight } = doc; if (scrollTop + clientHeight >= scrollHeight) { - page++; - if (props.details.md.length < props.details.total && page <= pageNum) { + if (props.details.md.length < props.details.total && page + 1 <= pageNum) { + page++; + loadingChunks = true; emit("getDoc", page); } } diff --git a/internal/application/service/agent_service.go b/internal/application/service/agent_service.go index 90219220..f7f50c28 100644 --- a/internal/application/service/agent_service.go +++ b/internal/application/service/agent_service.go @@ -109,7 +109,7 @@ func (s *agentService) CreateAgentEngine( // Register MCP tools from enabled services for this tenant tenantID := uint64(0) - if tid, ok := ctx.Value(types.TenantIDContextKey).(uint64); ok { + if tid, ok := types.TenantIDFromContext(ctx); ok { tenantID = tid } if tenantID > 0 && s.mcpServiceService != nil && s.mcpManager != nil { @@ -550,7 +550,7 @@ func (s *agentService) getSelectedDocumentInfos(ctx context.Context, knowledgeID // Get tenant ID from context tenantID := uint64(0) - if tid, ok := ctx.Value(types.TenantIDContextKey).(uint64); ok { + if tid, ok := types.TenantIDFromContext(ctx); ok { tenantID = tid } diff --git a/internal/application/service/chat_pipline/memory.go b/internal/application/service/chat_pipline/memory.go index 15bbf0f0..58ba2964 100644 --- a/internal/application/service/chat_pipline/memory.go +++ b/internal/application/service/chat_pipline/memory.go @@ -104,12 +104,12 @@ func (p *MemoryPlugin) handleStorage( {Role: "user", Content: chatManage.Query}, {Role: "assistant", Content: chatManage.ChatResponse.Content}, } - // Capture UserID and SessionID for goroutine userID := chatManage.UserID sessionID := chatManage.SessionID + bgCtx := context.WithoutCancel(ctx) go func() { - if err := p.memoryService.AddEpisode(ctx, userID, sessionID, messages); err != nil { - logger.Errorf(ctx, "failed to add episode: %v", err) + if err := p.memoryService.AddEpisode(bgCtx, userID, sessionID, messages); err != nil { + logger.Errorf(bgCtx, "failed to add episode: %v", err) } }() return nil @@ -118,11 +118,11 @@ func (p *MemoryPlugin) handleStorage( // If streaming, subscribe to events if chatManage.EventBus != nil { var fullResponse string - // Capture UserID and SessionID for closure userID := chatManage.UserID sessionID := chatManage.SessionID + bgCtx := context.WithoutCancel(ctx) - chatManage.EventBus.On(types.EventType(event.EventAgentFinalAnswer), func(ctx context.Context, evt types.Event) error { + chatManage.EventBus.On(types.EventType(event.EventAgentFinalAnswer), func(_ context.Context, evt types.Event) error { data, ok := evt.Data.(event.AgentFinalAnswerData) if !ok { return nil @@ -134,8 +134,8 @@ func (p *MemoryPlugin) handleStorage( {Role: "assistant", Content: fullResponse}, } go func() { - if err := p.memoryService.AddEpisode(ctx, userID, sessionID, messages); err != nil { - logger.Errorf(ctx, "failed to add episode: %v", err) + if err := p.memoryService.AddEpisode(bgCtx, userID, sessionID, messages); err != nil { + logger.Errorf(bgCtx, "failed to add episode: %v", err) } }() } diff --git a/internal/application/service/chat_pipline/merge.go b/internal/application/service/chat_pipline/merge.go index c1ed13ee..0ff3daca 100644 --- a/internal/application/service/chat_pipline/merge.go +++ b/internal/application/service/chat_pipline/merge.go @@ -220,7 +220,7 @@ func (p *PluginMerge) populateFAQAnswers( return results } - tenantID, _ := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, _ := types.TenantIDFromContext(ctx) if tenantID == 0 && chatManage != nil { tenantID = chatManage.TenantID } @@ -351,7 +351,7 @@ func (p *PluginMerge) expandShortContextWithNeighbors( return results } - tenantID, _ := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, _ := types.TenantIDFromContext(ctx) if tenantID == 0 && chatManage != nil { tenantID = chatManage.TenantID } diff --git a/internal/application/service/chat_pipline/search.go b/internal/application/service/chat_pipline/search.go index df55dd51..302ff414 100644 --- a/internal/application/service/chat_pipline/search.go +++ b/internal/application/service/chat_pipline/search.go @@ -503,7 +503,7 @@ func (p *PluginSearch) searchWebIfEnabled(ctx context.Context, chatManage *types if !chatManage.WebSearchEnabled || p.webSearchService == nil || p.tenantService == nil { return nil } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil || tenant.WebSearchConfig == nil || tenant.WebSearchConfig.Provider == "" { pipelineWarn(ctx, "Search", "web_config_missing", map[string]interface{}{ "tenant_id": chatManage.TenantID, diff --git a/internal/application/service/chat_pipline/search_entity.go b/internal/application/service/chat_pipline/search_entity.go index b5542527..a33c2b16 100644 --- a/internal/application/service/chat_pipline/search_entity.go +++ b/internal/application/service/chat_pipline/search_entity.go @@ -138,7 +138,7 @@ func (p *PluginSearchEntity) OnEvent(ctx context.Context, logger.Infof(ctx, "No new chunk found") return next() } - chunks, err := p.chunkRepo.ListChunksByID(ctx, ctx.Value(types.TenantIDContextKey).(uint64), chunkIDs) + chunks, err := p.chunkRepo.ListChunksByID(ctx, types.MustTenantIDFromContext(ctx), chunkIDs) if err != nil { logger.Errorf(ctx, "Failed to list chunks, session_id: %s, error: %v", chatManage.SessionID, err) return next() @@ -149,7 +149,7 @@ func (p *PluginSearchEntity) OnEvent(ctx context.Context, } knowledges, err := p.knowledgeRepo.GetKnowledgeBatch( ctx, - ctx.Value(types.TenantIDContextKey).(uint64), + types.MustTenantIDFromContext(ctx), knowledgeIDs, ) if err != nil { diff --git a/internal/application/service/chunk.go b/internal/application/service/chunk.go index b322c428..9533c70f 100644 --- a/internal/application/service/chunk.go +++ b/internal/application/service/chunk.go @@ -86,7 +86,7 @@ func (s *chunkService) CreateChunks(ctx context.Context, chunks []*types.Chunk) // - *types.Chunk: Retrieved chunk if found // - error: Any error encountered during retrieval func (s *chunkService) GetChunkByID(ctx context.Context, id string) (*types.Chunk, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Getting chunk by ID, ID: %s, tenant ID: %d", id, tenantID) chunk, err := s.chunkRepository.GetChunkByID(ctx, tenantID, id) if err != nil { @@ -126,7 +126,7 @@ func (s *chunkService) ListChunksByKnowledgeID(ctx context.Context, knowledgeID logger.Info(ctx, "Start listing chunks by knowledge ID") logger.Infof(ctx, "Knowledge ID: %s", knowledgeID) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) chunks, err := s.chunkRepository.ListChunksByKnowledgeID(ctx, tenantID, knowledgeID) @@ -155,7 +155,7 @@ func (s *chunkService) ListChunksByKnowledgeID(ctx context.Context, knowledgeID func (s *chunkService) ListPagedChunksByKnowledgeID(ctx context.Context, knowledgeID string, page *types.Pagination, chunkType []types.ChunkType, ) (*types.PageResult, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) chunks, total, err := s.chunkRepository.ListPagedChunksByKnowledgeID( ctx, tenantID, @@ -236,7 +236,7 @@ func (s *chunkService) UpdateChunks(ctx context.Context, chunks []*types.Chunk) // Returns: // - error: Any error encountered during deletion func (s *chunkService) DeleteChunk(ctx context.Context, id string) error { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) err := s.chunkRepository.DeleteChunk(ctx, tenantID, id) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ @@ -263,7 +263,7 @@ func (s *chunkService) DeleteChunks(ctx context.Context, ids []string) error { logger.Info(ctx, "Start deleting chunks in batch") logger.Infof(ctx, "Deleting %d chunks", len(ids)) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) err := s.chunkRepository.DeleteChunks(ctx, tenantID, ids) @@ -291,7 +291,7 @@ func (s *chunkService) DeleteChunksByKnowledgeID(ctx context.Context, knowledgeI logger.Info(ctx, "Start deleting all chunks by knowledge ID") logger.Infof(ctx, "Knowledge ID: %s", knowledgeID) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) err := s.chunkRepository.DeleteChunksByKnowledgeID(ctx, tenantID, knowledgeID) @@ -311,7 +311,7 @@ func (s *chunkService) DeleteByKnowledgeList(ctx context.Context, ids []string) logger.Info(ctx, "Start deleting all chunks by knowledge IDs") logger.Infof(ctx, "Knowledge IDs: %v", ids) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) err := s.chunkRepository.DeleteByKnowledgeList(ctx, tenantID, ids) @@ -353,7 +353,7 @@ func (s *chunkService) ListChunkByParentID( func (s *chunkService) DeleteGeneratedQuestion(ctx context.Context, chunkID string, questionID string) error { logger.Infof(ctx, "Deleting generated question, chunk ID: %s, question ID: %s", chunkID, questionID) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // 1. Get the chunk chunk, err := s.chunkRepository.GetChunkByID(ctx, tenantID, chunkID) @@ -404,7 +404,7 @@ func (s *chunkService) DeleteGeneratedQuestion(ctx context.Context, chunkID stri // The source_id format is: {chunk_id}-{question_id} sourceID := fmt.Sprintf("%s-%s", chunkID, questionID) - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenantInfo, _ := types.TenantInfoFromContext(ctx) retrieveEngine, err := retriever.NewCompositeRetrieveEngine(s.retrieveEngine, tenantInfo.GetEffectiveEngines()) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ diff --git a/internal/application/service/custom_agent.go b/internal/application/service/custom_agent.go index 65fb0b5c..433e73ba 100644 --- a/internal/application/service/custom_agent.go +++ b/internal/application/service/custom_agent.go @@ -46,7 +46,7 @@ func (s *customAgentService) CreateAgent(ctx context.Context, agent *types.Custo } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -90,7 +90,7 @@ func (s *customAgentService) GetAgentByID(ctx context.Context, id string) (*type } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -142,7 +142,7 @@ func (s *customAgentService) GetAgentByIDAndTenant(ctx context.Context, id strin // ListAgents lists all agents for the current tenant (including built-in agents) func (s *customAgentService) ListAgents(ctx context.Context) ([]*types.CustomAgent, error) { - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -204,7 +204,7 @@ func (s *customAgentService) UpdateAgent(ctx context.Context, agent *types.Custo } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } @@ -330,7 +330,7 @@ func (s *customAgentService) DeleteAgent(ctx context.Context, id string) error { } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return ErrInvalidTenantID } @@ -370,7 +370,7 @@ func (s *customAgentService) CopyAgent(ctx context.Context, id string) (*types.C } // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { return nil, ErrInvalidTenantID } diff --git a/internal/application/service/evaluation.go b/internal/application/service/evaluation.go index d6d1cb92..407f462e 100644 --- a/internal/application/service/evaluation.go +++ b/internal/application/service/evaluation.go @@ -109,7 +109,7 @@ func (e *EvaluationService) EvaluationResult(ctx context.Context, taskID string) return nil, err } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof( ctx, "Checking tenant ID match, task tenant ID: %d, current tenant ID: %d", @@ -138,7 +138,7 @@ func (e *EvaluationService) Evaluation(ctx context.Context, datasetID, knowledgeBaseID, chatModelID, rerankModelID) // Get tenant ID from context for multi-tenancy support - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) // Handle knowledge base creation if not provided diff --git a/internal/application/service/file/minio.go b/internal/application/service/file/minio.go index 0fcf2f42..2b3e617d 100644 --- a/internal/application/service/file/minio.go +++ b/internal/application/service/file/minio.go @@ -7,6 +7,7 @@ import ( "io" "mime/multipart" "path/filepath" + "strings" "time" "github.com/Tencent/WeKnora/internal/types/interfaces" @@ -89,6 +90,27 @@ func CheckMinioConnectivity(ctx context.Context, endpoint, accessKeyID, secretAc return svc.CheckConnectivity(ctx) } +// parseMinioFilePath extracts the object name from a "minio://bucket/key" path, +// validates the bucket matches this service instance, and checks for path traversal. +func (s *minioFileService) parseMinioFilePath(filePath string) (string, error) { + const prefix = "minio://" + if !strings.HasPrefix(filePath, prefix) { + return "", fmt.Errorf("invalid MinIO file path: %s", filePath) + } + rest := strings.TrimPrefix(filePath, prefix) + parts := strings.SplitN(rest, "/", 2) + if len(parts) != 2 || parts[0] == "" || parts[1] == "" { + return "", fmt.Errorf("invalid MinIO file path: %s", filePath) + } + if parts[0] != s.bucketName { + return "", fmt.Errorf("bucket mismatch in path: got %s, want %s", parts[0], s.bucketName) + } + if err := utils.SafeObjectKey(parts[1]); err != nil { + return "", fmt.Errorf("invalid file path: %w", err) + } + return parts[1], nil +} + // SaveFile saves a file to MinIO func (s *minioFileService) SaveFile(ctx context.Context, file *multipart.FileHeader, tenantID uint64, knowledgeID string, @@ -118,55 +140,28 @@ func (s *minioFileService) SaveFile(ctx context.Context, // GetFile gets a file from MinIO func (s *minioFileService) GetFile(ctx context.Context, filePath string) (io.ReadCloser, error) { - // Parse MinIO path - // Format: minio://bucketName/objectName - if len(filePath) < 9 || filePath[:8] != "minio://" { - return nil, fmt.Errorf("invalid MinIO file path: %s", filePath) + objectName, err := s.parseMinioFilePath(filePath) + if err != nil { + return nil, err } - - // Extract object name - objectName := filePath[9+len(s.bucketName):] - if len(objectName) > 0 && objectName[0] == '/' { - objectName = objectName[1:] - } - if err := utils.SafeObjectKey(objectName); err != nil { - return nil, fmt.Errorf("invalid file path: %w", err) - } - - // Get object obj, err := s.client.GetObject(ctx, s.bucketName, objectName, minio.GetObjectOptions{}) if err != nil { return nil, fmt.Errorf("failed to get file from MinIO: %w", err) } - return obj, nil } // DeleteFile deletes a file func (s *minioFileService) DeleteFile(ctx context.Context, filePath string) error { - // Parse MinIO path - // Format: minio://bucketName/objectName - if len(filePath) < 9 || filePath[:8] != "minio://" { - return fmt.Errorf("invalid MinIO file path: %s", filePath) - } - - // Extract object name - objectName := filePath[9+len(s.bucketName):] - if len(objectName) > 0 && objectName[0] == '/' { - objectName = objectName[1:] - } - if err := utils.SafeObjectKey(objectName); err != nil { - return fmt.Errorf("invalid file path: %w", err) - } - - // Delete object - err := s.client.RemoveObject(ctx, s.bucketName, objectName, minio.RemoveObjectOptions{ - GovernanceBypass: true, - }) + objectName, err := s.parseMinioFilePath(filePath) if err != nil { + return err + } + if err := s.client.RemoveObject(ctx, s.bucketName, objectName, minio.RemoveObjectOptions{ + GovernanceBypass: true, + }); err != nil { return fmt.Errorf("failed to delete file: %w", err) } - return nil } @@ -194,25 +189,13 @@ func (s *minioFileService) SaveBytes(ctx context.Context, data []byte, tenantID // GetFileURL returns a presigned download URL for the file func (s *minioFileService) GetFileURL(ctx context.Context, filePath string) (string, error) { - // Parse MinIO path - if len(filePath) < 9 || filePath[:8] != "minio://" { - return "", fmt.Errorf("invalid MinIO file path: %s", filePath) + objectName, err := s.parseMinioFilePath(filePath) + if err != nil { + return "", err } - - // Extract object name - objectName := filePath[9+len(s.bucketName):] - if len(objectName) > 0 && objectName[0] == '/' { - objectName = objectName[1:] - } - if err := utils.SafeObjectKey(objectName); err != nil { - return "", fmt.Errorf("invalid file path: %w", err) - } - - // Generate presigned URL (valid for 24 hours) presignedURL, err := s.client.PresignedGetObject(ctx, s.bucketName, objectName, 24*time.Hour, nil) if err != nil { return "", fmt.Errorf("failed to generate presigned URL: %w", err) } - return presignedURL.String(), nil } diff --git a/internal/application/service/image_multimodal.go b/internal/application/service/image_multimodal.go index 6fcad570..e5f3e1c6 100644 --- a/internal/application/service/image_multimodal.go +++ b/internal/application/service/image_multimodal.go @@ -4,8 +4,6 @@ import ( "context" "encoding/json" "fmt" - "io" - "net/http" "os" "strings" "time" @@ -16,6 +14,7 @@ import ( "github.com/Tencent/WeKnora/internal/models/vlm" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" + secutils "github.com/Tencent/WeKnora/internal/utils" "github.com/google/uuid" "github.com/hibiken/asynq" ) @@ -355,17 +354,5 @@ func (s *ImageMultimodalService) resolveVLM(ctx context.Context, kbID string) (v // downloadImageFromURL downloads image bytes from an HTTP(S) URL. func downloadImageFromURL(imageURL string) ([]byte, error) { - if !strings.HasPrefix(imageURL, "http://") && !strings.HasPrefix(imageURL, "https://") { - return nil, fmt.Errorf("unsupported URL scheme: %s", imageURL) - } - client := &http.Client{Timeout: 60 * time.Second} - resp, err := client.Get(imageURL) - if err != nil { - return nil, fmt.Errorf("HTTP GET: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, imageURL) - } - return io.ReadAll(resp.Body) + return secutils.DownloadBytes(imageURL) } diff --git a/internal/application/service/knowledgebase.go b/internal/application/service/knowledgebase.go index afd650a7..15f4eaac 100644 --- a/internal/application/service/knowledgebase.go +++ b/internal/application/service/knowledgebase.go @@ -82,7 +82,7 @@ func (s *knowledgeBaseService) CreateKnowledgeBase(ctx context.Context, kb.ID = uuid.New().String() } kb.CreatedAt = time.Now() - kb.TenantID = ctx.Value(types.TenantIDContextKey).(uint64) + kb.TenantID = types.MustTenantIDFromContext(ctx) kb.UpdatedAt = time.Now() kb.EnsureDefaults() @@ -158,7 +158,7 @@ func (s *knowledgeBaseService) GetKnowledgeBasesByIDsOnly(ctx context.Context, i // ListKnowledgeBases returns all knowledge bases for a tenant func (s *knowledgeBaseService) ListKnowledgeBases(ctx context.Context) ([]*types.KnowledgeBase, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) kbs, err := s.repo.ListKnowledgeBasesByTenantID(ctx, tenantID) if err != nil { @@ -324,8 +324,8 @@ func (s *knowledgeBaseService) DeleteKnowledgeBase(ctx context.Context, id strin logger.Infof(ctx, "Deleting knowledge base, ID: %s", id) // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenantID := types.MustTenantIDFromContext(ctx) + tenantInfo, _ := types.TenantInfoFromContext(ctx) // Step 1: Delete the knowledge base record first (mark as deleted) logger.Infof(ctx, "Deleting knowledge base from database") @@ -543,7 +543,7 @@ func (s *knowledgeBaseService) SetEmbeddingModel(ctx context.Context, id string, func (s *knowledgeBaseService) CopyKnowledgeBase(ctx context.Context, srcKB string, dstKB string, ) (*types.KnowledgeBase, *types.KnowledgeBase, error) { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Load source KB with tenant scope to prevent cross-tenant cloning sourceKB, err := s.repo.GetKnowledgeBaseByIDAndTenant(ctx, srcKB, tenantID) if err != nil { @@ -593,8 +593,8 @@ func (s *knowledgeBaseService) HybridSearch(ctx context.Context, ) ([]*types.SearchResult, error) { logger.Infof(ctx, "Hybrid search parameters, knowledge base ID: %s, query text: %s", id, params.QueryText) - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) - currentTenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantInfo, _ := types.TenantInfoFromContext(ctx) + currentTenantID := types.MustTenantIDFromContext(ctx) // Create a composite retrieval engine with tenant's configured retrievers retrieveEngine, err := retriever.NewCompositeRetrieveEngine(s.retrieveEngine, tenantInfo.GetEffectiveEngines()) @@ -880,7 +880,7 @@ func (s *knowledgeBaseService) iterativeRetrieveWithDeduplication(ctx context.Co filteredOutChunks := make(map[string]struct{}) queryTextLower := strings.ToLower(strings.TrimSpace(queryText)) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) for i := 0; i < maxIterations; i++ { // Update TopK in retrieve params @@ -1017,7 +1017,7 @@ func (s *knowledgeBaseService) filterByNegativeQuestions(ctx context.Context, return chunks } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Collect chunk IDs chunkIDs := make([]string, 0, len(chunks)) @@ -1104,7 +1104,7 @@ func (s *knowledgeBaseService) processSearchResults(ctx context.Context, return nil, nil } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Prepare data structures for efficient processing var knowledgeIDs []string diff --git a/internal/application/service/message.go b/internal/application/service/message.go index 1c5792d3..ea44c518 100644 --- a/internal/application/service/message.go +++ b/internal/application/service/message.go @@ -60,7 +60,7 @@ func (s *messageService) CreateMessage(ctx context.Context, message *types.Messa logger.Infof(ctx, "Creating message for session ID: %s", message.SessionID) // Check if the session exists to validate the message belongs to a valid session - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d, session ID: %s", tenantID, message.SessionID) _, err := s.sessionRepo.Get(ctx, tenantID, message.SessionID) if err != nil { @@ -94,7 +94,7 @@ func (s *messageService) GetMessage(ctx context.Context, sessionID string, messa logger.Infof(ctx, "Getting message, session ID: %s, message ID: %s", sessionID, messageID) // Verify the session exists before attempting to retrieve the message - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, sessionID) if err != nil { @@ -132,7 +132,7 @@ func (s *messageService) GetMessagesBySession(ctx context.Context, logger.Infof(ctx, "Getting messages for session ID: %s, page: %d, pageSize: %d", sessionID, page, pageSize) // Verify the session exists before retrieving messages - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, sessionID) if err != nil { @@ -253,7 +253,7 @@ func (s *messageService) UpdateMessage(ctx context.Context, message *types.Messa logger.Infof(ctx, "Updating message, ID: %s, session ID: %s", message.ID, message.SessionID) // Verify the session exists before updating the message - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, message.SessionID) if err != nil { @@ -288,7 +288,7 @@ func (s *messageService) DeleteMessage(ctx context.Context, sessionID string, me logger.Infof(ctx, "Deleting message, session ID: %s, message ID: %s", sessionID, messageID) // Verify the session exists before deleting the message - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Checking if session exists, tenant ID: %d", tenantID) _, err := s.sessionRepo.Get(ctx, tenantID, sessionID) if err != nil { diff --git a/internal/application/service/model.go b/internal/application/service/model.go index e4795224..dc553f9c 100644 --- a/internal/application/service/model.go +++ b/internal/application/service/model.go @@ -104,7 +104,7 @@ func (s *modelService) GetModelByID(ctx context.Context, id string) (*types.Mode return nil, errors.New("model ID cannot be empty") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Fetch model from repository model, err := s.repo.GetByID(ctx, tenantID, id) @@ -148,7 +148,7 @@ func (s *modelService) GetModelByID(ctx context.Context, id string) (*types.Mode func (s *modelService) ListModels(ctx context.Context) ([]*types.Model, error) { logger.Info(ctx, "Start listing models") - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Listing models for tenant ID: %d", tenantID) // List models from repository with no additional filters @@ -170,7 +170,7 @@ func (s *modelService) UpdateModel(ctx context.Context, model *types.Model) erro logger.Infof(ctx, "Updating model ID: %s, name: %s", model.ID, model.Name) // Check if the model is builtin - builtin models cannot be updated - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) existingModel, err := s.repo.GetByID(ctx, tenantID, model.ID) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ @@ -202,7 +202,7 @@ func (s *modelService) DeleteModel(ctx context.Context, id string) error { logger.Info(ctx, "Start deleting model") logger.Infof(ctx, "Deleting model ID: %s", id) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Tenant ID: %d", tenantID) // Check if the model is builtin - builtin models cannot be deleted @@ -368,7 +368,7 @@ func (s *modelService) GetChatModel(ctx context.Context, modelId string) (chat.C return nil, errors.New("model ID cannot be empty") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Get the model directly from repository to avoid status checks model, err := s.repo.GetByID(ctx, tenantID, modelId) @@ -412,7 +412,7 @@ func (s *modelService) GetVLMModel(ctx context.Context, modelId string) (vlm.VLM return nil, errors.New("model ID cannot be empty") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) model, err := s.repo.GetByID(ctx, tenantID, modelId) if err != nil { diff --git a/internal/application/service/session.go b/internal/application/service/session.go index ae670a4b..8a09f0f6 100644 --- a/internal/application/service/session.go +++ b/internal/application/service/session.go @@ -114,7 +114,7 @@ func (s *sessionService) GetSession(ctx context.Context, id string) (*types.Sess } // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Retrieving session, ID: %s, tenant ID: %d", id, tenantID) // Get session from repository @@ -134,7 +134,7 @@ func (s *sessionService) GetSession(ctx context.Context, id string) (*types.Sess // GetSessionsByTenant retrieves all sessions for the current tenant func (s *sessionService) GetSessionsByTenant(ctx context.Context) ([]*types.Session, error) { // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) logger.Infof(ctx, "Retrieving all sessions for tenant, tenant ID: %d", tenantID) // Get sessions from repository @@ -157,7 +157,7 @@ func (s *sessionService) GetPagedSessionsByTenant(ctx context.Context, pagination *types.Pagination, ) (*types.PageResult, error) { // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Get paged sessions from repository sessions, total, err := s.sessionRepo.GetPagedByTenantID(ctx, tenantID, pagination) if err != nil { @@ -203,7 +203,7 @@ func (s *sessionService) DeleteSession(ctx context.Context, id string) error { } // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Cleanup temporary KB stored in Redis for this session if err := s.webSearchStateRepo.DeleteWebSearchTempKBState(ctx, id); err != nil { @@ -236,7 +236,7 @@ func (s *sessionService) BatchDeleteSessions(ctx context.Context, ids []string) } // Get tenant ID from context - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) // Cleanup associated resources for each session for _, id := range ids { @@ -643,7 +643,7 @@ func (s *sessionService) KnowledgeQA( ) // Get UserID from context - userID, _ := ctx.Value(types.UserIDContextKey).(string) + userID, _ := types.UserIDFromContext(ctx) chatManage := &types.ChatManage{ Query: query, @@ -790,7 +790,7 @@ func (s *sessionService) selectChatModelID( ) (string, error) { // If no knowledge base IDs but have knowledge IDs, derive KB IDs from knowledge IDs (include shared KB files) if len(knowledgeBaseIDs) == 0 && len(knowledgeIDs) > 0 { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) knowledgeList, err := s.knowledgeService.GetKnowledgeBatchWithSharedAccess(ctx, tenantID, knowledgeIDs) if err != nil { logger.Warnf(ctx, "Failed to get knowledge batch for model selection: %v", err) @@ -899,7 +899,7 @@ func (s *sessionService) resolveKnowledgeBasesFromAgent( // unrelated KBs from other organisations into the agent's retrieval scope. isSharedAgent := sessionTenantID != 0 && sessionTenantID != customAgent.TenantID if !isSharedAgent { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) userIDVal := ctx.Value(types.UserIDContextKey) if userIDVal != nil { if userID, ok := userIDVal.(string); ok && userID != "" && s.kbShareService != nil { @@ -1021,7 +1021,7 @@ func (s *sessionService) buildSearchTargets( kbByID[kb.ID] = kb } } - userID, _ := ctx.Value(types.UserIDContextKey).(string) + userID, _ := types.UserIDFromContext(ctx) for _, kbID := range knowledgeBaseIDs { fullKBSet[kbID] = true kb := kbByID[kbID] @@ -1114,7 +1114,7 @@ func (s *sessionService) KnowledgeQAByEvent(ctx context.Context, // Set up tracing attributes logger.Infof(ctx, "Trigger event list: %v", methods) span.SetAttributes( - attribute.String("request_id", ctx.Value(types.RequestIDContextKey).(string)), + attribute.String("request_id", func() string { id, _ := types.RequestIDFromContext(ctx); return id }()), attribute.String("query", chatManage.Query), attribute.String("method", strings.Join(methods, ",")), ) @@ -1163,7 +1163,7 @@ func (s *sessionService) SearchKnowledge(ctx context.Context, knowledgeBaseIDs, knowledgeIDs, query) // Get tenant ID from context - tenantID, ok := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID, ok := types.TenantIDFromContext(ctx) if !ok { logger.Error(ctx, "Failed to get tenant ID from context") return nil, fmt.Errorf("tenant ID not found in context") @@ -1181,7 +1181,7 @@ func (s *sessionService) SearchKnowledge(ctx context.Context, } // Create default retrieval parameters - userID, _ := ctx.Value(types.UserIDContextKey).(string) + userID, _ := types.UserIDFromContext(ctx) chatManage := &types.ChatManage{ Query: query, RewriteQuery: query, @@ -1528,7 +1528,7 @@ func (s *sessionService) getContextManagerForSession( chatModel chat.Chat, ) interfaces.ContextManager { // Get tenant to access global context configuration - tenant, _ := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) // Determine which context config to use: tenant-level or default var contextConfig *types.ContextConfig if tenant != nil && tenant.ContextConfig != nil { diff --git a/internal/application/service/tag.go b/internal/application/service/tag.go index 51bb748f..7e9fb88a 100644 --- a/internal/application/service/tag.go +++ b/internal/application/service/tag.go @@ -73,7 +73,7 @@ func (s *knowledgeTagService) ListTags( } // Check access permission - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) if kb.TenantID != tenantID { // Get user ID from context userIDVal := ctx.Value(types.UserIDContextKey) @@ -192,7 +192,7 @@ func (s *knowledgeTagService) UpdateTag( if id == "" { return nil, werrors.NewBadRequestError("标签ID不能为空") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) tag, err := s.repo.GetByID(ctx, tenantID, id) if err != nil { return nil, err @@ -225,7 +225,7 @@ func (s *knowledgeTagService) DeleteTag(ctx context.Context, id string, force bo if id == "" { return werrors.NewBadRequestError("标签ID不能为空") } - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) tag, err := s.repo.GetByID(ctx, tenantID, id) if err != nil { return err @@ -243,7 +243,7 @@ func (s *knowledgeTagService) DeleteTag(ctx context.Context, id string, force bo } // Get tenant info for effective engines - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenantInfo, _ := types.TenantInfoFromContext(ctx) // Helper function to delete chunks and enqueue index deletion task deleteChunksAndEnqueueIndexDelete := func() error { diff --git a/internal/handler/initialization.go b/internal/handler/initialization.go index e1d3395d..094d6fca 100644 --- a/internal/handler/initialization.go +++ b/internal/handler/initialization.go @@ -1912,7 +1912,7 @@ func (h *InitializationHandler) testMultimodalWithDocReader( return nil, fmt.Errorf("DocReader service not configured") } - requestID, _ := ctx.Value(types.RequestIDContextKey).(string) + requestID, _ := types.RequestIDFromContext(ctx) readResult, err := h.documentReader.Read(ctx, &types.ReadRequest{ FileContent: imageContent, diff --git a/internal/handler/knowledge.go b/internal/handler/knowledge.go index b68fcfac..c6e582b7 100644 --- a/internal/handler/knowledge.go +++ b/internal/handler/knowledge.go @@ -44,74 +44,10 @@ func NewKnowledgeHandler( } // validateKnowledgeBaseAccess validates access permissions to a knowledge base -// Returns the knowledge base, the knowledge base ID, effective tenant ID, permission level, and any errors encountered -// For owned KBs, effectiveTenantID is the caller's tenant ID -// For shared KBs, effectiveTenantID is the source tenant ID (owner's tenant) +// using the ":id" URL path parameter. It delegates to validateKnowledgeBaseAccessWithKBID. func (h *KnowledgeHandler) validateKnowledgeBaseAccess(c *gin.Context) (*types.KnowledgeBase, string, uint64, types.OrgMemberRole, error) { - ctx := c.Request.Context() - - // Get tenant ID from context - tenantID := c.GetUint64(types.TenantIDContextKey.String()) - if tenantID == 0 { - logger.Error(ctx, "Failed to get tenant ID") - return nil, "", 0, "", errors.NewUnauthorizedError("Unauthorized") - } - - // Get user ID from context (needed for shared KB permission check) - userID, userExists := c.Get(types.UserIDContextKey.String()) - - // Get knowledge base ID from URL path parameter kbID := secutils.SanitizeForLog(c.Param("id")) - if kbID == "" { - logger.Error(ctx, "Knowledge base ID is empty") - return nil, "", 0, "", errors.NewBadRequestError("Knowledge base ID cannot be empty") - } - - // Get knowledge base details - kb, err := h.kbService.GetKnowledgeBaseByID(ctx, kbID) - if err != nil { - logger.ErrorWithFields(ctx, err, nil) - return nil, kbID, 0, "", errors.NewInternalServerError(err.Error()) - } - - // Check 1: Verify tenant ownership (owner has full access) - if kb.TenantID == tenantID { - return kb, kbID, tenantID, types.OrgRoleAdmin, nil - } - - // Check 2: If not owner, check organization shared access - if userExists && h.kbShareService != nil { - // Check if user has shared access through organization - permission, isShared, permErr := h.kbShareService.CheckUserKBPermission(ctx, kbID, userID.(string)) - if permErr == nil && isShared { - // User has shared access, get the source tenant ID for queries - sourceTenantID, srcErr := h.kbShareService.GetKBSourceTenant(ctx, kbID) - if srcErr == nil { - logger.Infof(ctx, "User %s accessing shared KB %s with permission %s, source tenant: %d", - userID.(string), kbID, permission, sourceTenantID) - return kb, kbID, sourceTenantID, permission, nil - } - } - } - - // Check 3: If not owner and no direct share, allow if user has any shared agent that can access this KB (e.g. opened from "通过智能体可见" list) - if userExists && h.agentShareService != nil { - can, err := h.agentShareService.UserCanAccessKBViaSomeSharedAgent(ctx, userID.(string), tenantID, kb) - if err == nil && can { - logger.Infof(ctx, "User %s accessing KB %s via some shared agent", userID.(string), kbID) - return kb, kbID, kb.TenantID, types.OrgRoleViewer, nil - } - } - - // No permission: not owner and no shared access - logger.Warnf( - ctx, - "Permission denied to access this knowledge base, tenant ID mismatch, "+ - "requested tenant ID: %d, knowledge base tenant ID: %d", - tenantID, - kb.TenantID, - ) - return nil, kbID, 0, "", errors.NewForbiddenError("Permission denied to access this knowledge base") + return h.validateKnowledgeBaseAccessWithKBID(c, kbID) } // validateKnowledgeBaseAccessWithKBID validates access to the given knowledge base ID (e.g. from query or body). diff --git a/internal/handler/organization.go b/internal/handler/organization.go index 8d24d1e5..a4286819 100644 --- a/internal/handler/organization.go +++ b/internal/handler/organization.go @@ -1225,7 +1225,7 @@ func (h *OrganizationHandler) ListSharedKnowledgeBases(c *gin.Context) { ctx := c.Request.Context() userID := c.GetString(types.UserIDContextKey.String()) - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) sharedKBs, err := h.shareService.ListSharedKnowledgeBases(ctx, userID, tenantID) if err != nil { diff --git a/internal/handler/session/helpers.go b/internal/handler/session/helpers.go index e9c2da34..03abf522 100644 --- a/internal/handler/session/helpers.go +++ b/internal/handler/session/helpers.go @@ -215,7 +215,7 @@ func getFloat64(m map[string]interface{}, key string) float64 { // It prioritizes tenant-level ConversationConfig, then falls back to config.yaml defaults func (h *Handler) createDefaultSummaryConfig(ctx context.Context) *types.SummaryConfig { // Try to get tenant from context - tenant, _ := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) // Initialize with config.yaml defaults cfg := &types.SummaryConfig{ @@ -259,7 +259,7 @@ func (h *Handler) createDefaultSummaryConfig(ctx context.Context) *types.Summary // It prioritizes tenant-level ConversationConfig, then falls back to config.yaml defaults func (h *Handler) fillSummaryConfigDefaults(ctx context.Context, config *types.SummaryConfig) { // Try to get tenant from context - tenant, _ := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) // Determine default values: tenant config first, then config.yaml var defaultPrompt, defaultContextTemplate, defaultNoMatchPrefix string diff --git a/internal/handler/tag.go b/internal/handler/tag.go index 7db819de..25b3e7b2 100644 --- a/internal/handler/tag.go +++ b/internal/handler/tag.go @@ -94,7 +94,7 @@ func (h *TagHandler) resolveTagIDWithCtx(c *gin.Context, ctx context.Context) (s tagIDParam := secutils.SanitizeForLog(c.Param("tag_id")) if seqID, err := strconv.ParseInt(tagIDParam, 10, 64); err == nil { - tenantID := ctx.Value(types.TenantIDContextKey).(uint64) + tenantID := types.MustTenantIDFromContext(ctx) tag, err := h.tagRepo.GetBySeqID(ctx, tenantID, seqID) if err != nil { return "", errors.NewNotFoundError("标签不存在") diff --git a/internal/handler/tenant.go b/internal/handler/tenant.go index 2e7db31f..a89e2e6d 100644 --- a/internal/handler/tenant.go +++ b/internal/handler/tenant.go @@ -439,7 +439,7 @@ type AgentConfigRequest struct { // @Router /tenants/kv/agent-config [get] func (h *TenantHandler) GetTenantAgentConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -529,7 +529,7 @@ func (h *TenantHandler) updateTenantAgentConfigInternal(c *gin.Context) { } // Get existing tenant - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -670,7 +670,7 @@ func (h *TenantHandler) updateTenantWebSearchConfigInternal(c *gin.Context) { return } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -711,7 +711,7 @@ func (h *TenantHandler) GetTenantWebSearchConfig(c *gin.Context) { ctx := c.Request.Context() logger.Info(ctx, "Start getting tenant web search config") // Get tenant - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -728,7 +728,7 @@ func (h *TenantHandler) GetTenantWebSearchConfig(c *gin.Context) { // GetTenantParserEngineConfig returns the tenant's parser engine config (MinerU endpoint, API key, etc.). func (h *TenantHandler) GetTenantParserEngineConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -753,7 +753,7 @@ func (h *TenantHandler) updateTenantParserEngineConfigInternal(c *gin.Context) { c.Error(errors.NewValidationError("Invalid request data").WithDetails(err.Error())) return } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -780,7 +780,7 @@ func (h *TenantHandler) updateTenantParserEngineConfigInternal(c *gin.Context) { // GetTenantStorageEngineConfig returns the tenant's storage engine config (Local, MinIO, COS parameters). func (h *TenantHandler) GetTenantStorageEngineConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -805,7 +805,7 @@ func (h *TenantHandler) updateTenantStorageEngineConfigInternal(c *gin.Context) c.Error(errors.NewValidationError("Invalid request data").WithDetails(err.Error())) return } - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -897,7 +897,7 @@ func validateConversationConfig(req *types.ConversationConfig) error { // @Router /tenants/kv/conversation-config [get] func (h *TenantHandler) GetTenantConversationConfig(c *gin.Context) { ctx := c.Request.Context() - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) @@ -934,7 +934,7 @@ func (h *TenantHandler) updateTenantConversationInternal(c *gin.Context) { } // Get existing tenant - tenant := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + tenant, _ := types.TenantInfoFromContext(ctx) if tenant == nil { logger.Error(ctx, "Tenant is empty") c.Error(errors.NewBadRequestError("Tenant is empty")) diff --git a/internal/infrastructure/docparser/helpers.go b/internal/infrastructure/docparser/helpers.go new file mode 100644 index 00000000..b0ba4d9b --- /dev/null +++ b/internal/infrastructure/docparser/helpers.go @@ -0,0 +1,110 @@ +package docparser + +import ( + "context" + "fmt" + "sort" + "strings" + "time" +) + +// stringOr returns val (trimmed) if non-empty, otherwise fallback. +func stringOr(val, fallback string) string { + val = strings.TrimSpace(val) + if val == "" { + return fallback + } + return val +} + +// parseBoolOr parses a truthy string ("true","1","yes"), returning fallback on empty. +func parseBoolOr(val string, fallback bool) bool { + val = strings.ToLower(strings.TrimSpace(val)) + if val == "" { + return fallback + } + return val == "true" || val == "1" || val == "yes" +} + +// firstNonEmpty returns the first non-empty string, or "" if all are empty. +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} + +// sleepCtx sleeps for d but returns early if ctx is cancelled. +func sleepCtx(ctx context.Context, d time.Duration) { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + case <-t.C: + } +} + +// logResponseStructure recursively logs the structure of an API response, +// truncating large string values. label identifies the subsystem (e.g. "MinerU"). +func logResponseStructure(label string, obj interface{}, prefix string) { + switch v := obj.(type) { + case map[string]interface{}: + keys := make([]string, 0, len(v)) + for k := range v { + keys = append(keys, k) + } + sort.Strings(keys) + logger.Printf("DEBUG: [%s] %s = {object with %d keys: %s}", label, prefix, len(v), strings.Join(keys, ", ")) + for _, key := range keys { + val := v[key] + path := prefix + "." + key + switch inner := val.(type) { + case map[string]interface{}: + logResponseStructure(label, inner, path) + case []interface{}: + logger.Printf("DEBUG: [%s] %s = [array with %d items]", label, path, len(inner)) + if len(inner) > 0 { + logger.Printf("DEBUG: [%s] %s[0] type=%T", label, path, inner[0]) + if len(inner) <= 3 { + for i, item := range inner { + logResponseStructure(label, item, fmt.Sprintf("%s[%d]", path, i)) + } + } else { + logResponseStructure(label, inner[0], path+"[0]") + logger.Printf("DEBUG: [%s] ... and %d more items in %s", label, len(inner)-1, path) + } + } + case string: + if len(inner) > 200 { + logger.Printf("DEBUG: [%s] %s = string(%d chars): %.200s...", label, path, len(inner), inner) + } else { + logger.Printf("DEBUG: [%s] %s = %q", label, path, inner) + } + case float64: + logger.Printf("DEBUG: [%s] %s = %v (number)", label, path, inner) + case bool: + logger.Printf("DEBUG: [%s] %s = %v (bool)", label, path, inner) + case nil: + logger.Printf("DEBUG: [%s] %s = null", label, path) + default: + logger.Printf("DEBUG: [%s] %s = %v (%T)", label, path, val, val) + } + } + case []interface{}: + logger.Printf("DEBUG: [%s] %s = [array with %d items]", label, prefix, len(v)) + if len(v) > 0 { + if len(v) <= 3 { + for i, item := range v { + logResponseStructure(label, item, fmt.Sprintf("%s[%d]", prefix, i)) + } + } else { + logResponseStructure(label, v[0], prefix+"[0]") + logger.Printf("DEBUG: [%s] ... and %d more items in %s", label, len(v)-1, prefix) + } + } + default: + logger.Printf("DEBUG: [%s] %s = %v (%T)", label, prefix, v, v) + } +} diff --git a/internal/infrastructure/docparser/image_resolver.go b/internal/infrastructure/docparser/image_resolver.go index a8a54eda..964a9349 100644 --- a/internal/infrastructure/docparser/image_resolver.go +++ b/internal/infrastructure/docparser/image_resolver.go @@ -2,20 +2,16 @@ package docparser import ( "fmt" - "io" - "net/http" "os" "path/filepath" "regexp" "strings" - "time" "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/utils" "github.com/google/uuid" ) -var httpImageClient = &http.Client{Timeout: 60 * time.Second} - // StoredImage describes an image that has been saved to local storage. type StoredImage struct { OriginalRef string // reference in the original markdown @@ -179,19 +175,7 @@ func (r *ImageResolver) readImageBytes( } func downloadImage(url string) ([]byte, error) { - resp, err := httpImageClient.Get(url) - if err != nil { - return nil, fmt.Errorf("HTTP GET: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url) - } - data, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("read body: %w", err) - } - return data, nil + return utils.DownloadBytes(url) } func extFromMime(mime string) string { diff --git a/internal/infrastructure/docparser/mineru_cloud_converter.go b/internal/infrastructure/docparser/mineru_cloud_converter.go index 44bd08c1..4171684e 100644 --- a/internal/infrastructure/docparser/mineru_cloud_converter.go +++ b/internal/infrastructure/docparser/mineru_cloud_converter.go @@ -297,7 +297,7 @@ func (c *MinerUCloudReader) fetchBatchStatus(ctx context.Context, batchID string // Pretty-print the structure to reveal all available fields var rawObj interface{} if err := json.Unmarshal(pollResp.Data.ExtractResult, &rawObj); err == nil { - logCloudResponseStructure(rawObj, "extract_result") + logResponseStructure("MinerUCloud", rawObj, "extract_result") } // The extract_result can be either a single object or an array @@ -472,68 +472,7 @@ func readZipEntryBytes(f *zip.File) ([]byte, error) { return io.ReadAll(rc) } -// logCloudResponseStructure recursively logs the structure of the Cloud API response, -// truncating large string values to reveal all available fields (e.g. caption, bbox, etc.). -func logCloudResponseStructure(obj interface{}, prefix string) { - switch v := obj.(type) { - case map[string]interface{}: - keys := make([]string, 0, len(v)) - for k := range v { - keys = append(keys, k) - } - sort.Strings(keys) - logger.Printf("DEBUG: [MinerUCloud] %s = {object with %d keys: %s}", prefix, len(v), strings.Join(keys, ", ")) - for _, key := range keys { - val := v[key] - path := prefix + "." + key - switch inner := val.(type) { - case map[string]interface{}: - logCloudResponseStructure(inner, path) - case []interface{}: - logger.Printf("DEBUG: [MinerUCloud] %s = [array with %d items]", path, len(inner)) - if len(inner) > 0 { - logger.Printf("DEBUG: [MinerUCloud] %s[0] type=%T", path, inner[0]) - if len(inner) <= 3 { - for i, item := range inner { - logCloudResponseStructure(item, fmt.Sprintf("%s[%d]", path, i)) - } - } else { - logCloudResponseStructure(inner[0], path+"[0]") - logger.Printf("DEBUG: [MinerUCloud] ... and %d more items in %s", len(inner)-1, path) - } - } - case string: - if len(inner) > 200 { - logger.Printf("DEBUG: [MinerUCloud] %s = string(%d chars): %.200s...", path, len(inner), inner) - } else { - logger.Printf("DEBUG: [MinerUCloud] %s = %q", path, inner) - } - case float64: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (number)", path, inner) - case bool: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (bool)", path, inner) - case nil: - logger.Printf("DEBUG: [MinerUCloud] %s = null", path) - default: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (%T)", path, val, val) - } - } - case []interface{}: - logger.Printf("DEBUG: [MinerUCloud] %s = [array with %d items]", prefix, len(v)) - if len(v) > 0 { - if len(v) <= 3 { - for i, item := range v { - logCloudResponseStructure(item, fmt.Sprintf("%s[%d]", prefix, i)) - } - } else { - logCloudResponseStructure(v[0], prefix+"[0]") - logger.Printf("DEBUG: [MinerUCloud] ... and %d more items in %s", len(v)-1, prefix) - } - } - default: - logger.Printf("DEBUG: [MinerUCloud] %s = %v (%T)", prefix, v, v) - } -} +// logCloudResponseStructure is now unified in helpers.go as logResponseStructure("MinerUCloud", ...) // PingMinerUCloud checks if the MinerU Cloud API is reachable with the given API key. func PingMinerUCloud(apiKey, baseURL string) (bool, string) { @@ -568,22 +507,4 @@ func PingMinerUCloud(apiKey, baseURL string) (bool, string) { return true, "" } -// --- helpers --- - -func firstNonEmpty(vals ...string) string { - for _, v := range vals { - if v != "" { - return v - } - } - return "" -} - -func sleepCtx(ctx context.Context, d time.Duration) { - t := time.NewTimer(d) - defer t.Stop() - select { - case <-ctx.Done(): - case <-t.C: - } -} +// firstNonEmpty, sleepCtx are defined in helpers.go diff --git a/internal/infrastructure/docparser/mineru_converter.go b/internal/infrastructure/docparser/mineru_converter.go index a70fef0c..11394fb1 100644 --- a/internal/infrastructure/docparser/mineru_converter.go +++ b/internal/infrastructure/docparser/mineru_converter.go @@ -25,12 +25,12 @@ var b64DataURIPattern = regexp.MustCompile(`^data:image/(\w+);base64,(.+)$`) // MinerUReader calls a self-hosted MinerU API to read/convert documents. type MinerUReader struct { - endpoint string - backend string // "pipeline", "vlm-*", "hybrid-*" - formulaEnable bool - tableEnable bool - ocrEnable bool - language string + endpoint string + backend string // "pipeline", "vlm-*", "hybrid-*" + formulaEnable bool + tableEnable bool + ocrEnable bool + language string } // NewMinerUReader creates a reader from ParserEngineOverrides. @@ -95,18 +95,18 @@ func (c *MinerUReader) callFileParse(ctx context.Context, content []byte) (strin // Form fields fields := map[string]string{ - "return_md": "true", - "return_images": "true", - "table_enable": fmt.Sprintf("%v", c.tableEnable), - "formula_enable": fmt.Sprintf("%v", c.formulaEnable), - "parse_method": "ocr", - "start_page_id": "0", - "end_page_id": "99999", - "backend": c.backend, - "response_format_zip": "false", - "return_middle_json": "false", - "return_model_output": "false", - "return_content_list": "true", + "return_md": "true", + "return_images": "true", + "table_enable": fmt.Sprintf("%v", c.tableEnable), + "formula_enable": fmt.Sprintf("%v", c.formulaEnable), + "parse_method": "ocr", + "start_page_id": "0", + "end_page_id": "99999", + "backend": c.backend, + "response_format_zip": "false", + "return_middle_json": "false", + "return_model_output": "false", + "return_content_list": "true", } if !c.ocrEnable { fields["parse_method"] = "txt" @@ -162,7 +162,7 @@ func (c *MinerUReader) callFileParse(ctx context.Context, content []byte) (strin // Also pretty-print the top-level structure (without large base64 blobs) var rawMap map[string]interface{} if err := json.Unmarshal(respBody, &rawMap); err == nil { - c.logResponseStructure(rawMap, "") + c.logMinerUResponseStructure(rawMap, "") } var result mineruFileParseResponse @@ -225,60 +225,9 @@ func (c *MinerUReader) processImages(mdContent string, imagesB64 map[string]stri return refs, mdContent } -// logResponseStructure recursively logs the structure of the API response, -// truncating large string values to reveal all available fields. -func (c *MinerUReader) logResponseStructure(obj interface{}, prefix string) { - switch v := obj.(type) { - case map[string]interface{}: - for key, val := range v { - path := prefix + "." + key - if path[0] == '.' { - path = path[1:] - } - switch inner := val.(type) { - case map[string]interface{}: - logger.Printf("DEBUG: [MinerU] Response field: %s = {object with %d keys: %s}", path, len(inner), mapKeys(inner)) - c.logResponseStructure(inner, path) - case []interface{}: - logger.Printf("DEBUG: [MinerU] Response field: %s = [array with %d items]", path, len(inner)) - if len(inner) > 0 { - logger.Printf("DEBUG: [MinerU] Response field: %s[0] type=%T", path, inner[0]) - if len(inner) <= 3 { - for i, item := range inner { - c.logResponseStructure(item, fmt.Sprintf("%s[%d]", path, i)) - } - } else { - c.logResponseStructure(inner[0], path+"[0]") - logger.Printf("DEBUG: [MinerU] ... and %d more items in %s", len(inner)-1, path) - } - } - case string: - if len(inner) > 200 { - logger.Printf("DEBUG: [MinerU] Response field: %s = string(%d chars): %.200s...", path, len(inner), inner) - } else { - logger.Printf("DEBUG: [MinerU] Response field: %s = %q", path, inner) - } - case float64: - logger.Printf("DEBUG: [MinerU] Response field: %s = %v (number)", path, inner) - case bool: - logger.Printf("DEBUG: [MinerU] Response field: %s = %v (bool)", path, inner) - case nil: - logger.Printf("DEBUG: [MinerU] Response field: %s = null", path) - default: - logger.Printf("DEBUG: [MinerU] Response field: %s = %v (%T)", path, inner, inner) - } - } - default: - logger.Printf("DEBUG: [MinerU] Value at %s: %v (%T)", prefix, v, v) - } -} - -func mapKeys(m map[string]interface{}) string { - keys := make([]string, 0, len(m)) - for k := range m { - keys = append(keys, k) - } - return strings.Join(keys, ", ") +// logMinerUResponseStructure logs the structure of the MinerU API response. +func (c *MinerUReader) logMinerUResponseStructure(obj interface{}, prefix string) { + logResponseStructure("MinerU", obj, prefix) } // PingMinerU checks if the self-hosted MinerU service is reachable. @@ -309,21 +258,3 @@ func htmlToMarkdown(content string) string { } return md } - -// --- helpers --- - -func stringOr(val, fallback string) string { - val = strings.TrimSpace(val) - if val == "" { - return fallback - } - return val -} - -func parseBoolOr(val string, fallback bool) bool { - val = strings.ToLower(strings.TrimSpace(val)) - if val == "" { - return fallback - } - return val == "true" || val == "1" || val == "yes" -} diff --git a/internal/types/context_helpers.go b/internal/types/context_helpers.go new file mode 100644 index 00000000..0f383787 --- /dev/null +++ b/internal/types/context_helpers.go @@ -0,0 +1,47 @@ +package types + +import "context" + +// TenantIDFromContext extracts the tenant ID from ctx. +// Returns (0, false) when the key is absent or the value is not uint64. +func TenantIDFromContext(ctx context.Context) (uint64, bool) { + v, ok := ctx.Value(TenantIDContextKey).(uint64) + return v, ok +} + +// MustTenantIDFromContext extracts the tenant ID from ctx, panicking if missing. +func MustTenantIDFromContext(ctx context.Context) uint64 { + v, ok := TenantIDFromContext(ctx) + if !ok { + panic("types.TenantIDContextKey not set in context") + } + return v +} + +// TenantInfoFromContext extracts the *Tenant from ctx. +func TenantInfoFromContext(ctx context.Context) (*Tenant, bool) { + v, ok := ctx.Value(TenantInfoContextKey).(*Tenant) + return v, ok && v != nil +} + +// RequestIDFromContext extracts the request ID string from ctx. +func RequestIDFromContext(ctx context.Context) (string, bool) { + v, ok := ctx.Value(RequestIDContextKey).(string) + return v, ok && v != "" +} + +// UserIDFromContext extracts the user ID string from ctx. +func UserIDFromContext(ctx context.Context) (string, bool) { + v, ok := ctx.Value(UserIDContextKey).(string) + return v, ok && v != "" +} + +// SessionTenantIDFromContext extracts the session-owner tenant ID from ctx. +// Falls back to TenantIDFromContext when the session key is absent. +func SessionTenantIDFromContext(ctx context.Context) (uint64, bool) { + v, ok := ctx.Value(SessionTenantIDContextKey).(uint64) + if ok && v != 0 { + return v, true + } + return TenantIDFromContext(ctx) +} diff --git a/internal/utils/httputil.go b/internal/utils/httputil.go new file mode 100644 index 00000000..b1d3ed0a --- /dev/null +++ b/internal/utils/httputil.go @@ -0,0 +1,32 @@ +package utils + +import ( + "fmt" + "io" + "net/http" + "strings" + "time" +) + +var defaultHTTPClient = &http.Client{Timeout: 60 * time.Second} + +// DownloadBytes fetches the content at the given HTTP(S) URL and returns the +// raw bytes. It reuses a package-level http.Client with a 60-second timeout. +func DownloadBytes(url string) ([]byte, error) { + if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") { + return nil, fmt.Errorf("unsupported URL scheme: %s", url) + } + resp, err := defaultHTTPClient.Get(url) + if err != nil { + return nil, fmt.Errorf("HTTP GET: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("HTTP %d for %s", resp.StatusCode, url) + } + data, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read body: %w", err) + } + return data, nil +}