Files
Enterprise-Knowledge-Base/internal/application/service/custom_agent.go
T
wizardchen 4b4225a79f feat(custom_agent): enhance custom agent service with KB sharing support
Updated the customAgentService to include a new KBShareService interface for handling cross-tenant knowledge base sharing. Implemented logic to group knowledge base IDs by effective tenant, ensuring correct querying of shared KBs. Adjusted the GetSuggestedQuestions method to accommodate these changes, improving the handling of FAQ and document chunks for shared knowledge bases.

Refs: #1303
2026-05-18 17:28:58 +08:00

807 lines
24 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package service
import (
"context"
"errors"
"math/rand"
"strings"
"time"
"github.com/Tencent/WeKnora/internal/agent/tools"
"github.com/Tencent/WeKnora/internal/application/repository"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/google/uuid"
)
// Custom agent related errors
var (
ErrAgentNotFound = errors.New("agent not found")
ErrCannotModifyBuiltin = errors.New("cannot modify built-in agent basic info")
ErrCannotDeleteBuiltin = errors.New("cannot delete built-in agent")
ErrAgentNameRequired = errors.New("agent name is required")
)
// customAgentService implements the CustomAgentService interface
type customAgentService struct {
repo interfaces.CustomAgentRepository
chunkRepo interfaces.ChunkRepository
kbService interfaces.KnowledgeBaseService
kbShareService interfaces.KBShareService
wikiPageRepo interfaces.WikiPageRepository
}
// NewCustomAgentService creates a new custom agent service
func NewCustomAgentService(
repo interfaces.CustomAgentRepository,
chunkRepo interfaces.ChunkRepository,
kbService interfaces.KnowledgeBaseService,
kbShareService interfaces.KBShareService,
wikiPageRepo interfaces.WikiPageRepository,
) interfaces.CustomAgentService {
return &customAgentService{
repo: repo,
chunkRepo: chunkRepo,
kbService: kbService,
kbShareService: kbShareService,
wikiPageRepo: wikiPageRepo,
}
}
// CreateAgent creates a new custom agent
func (s *customAgentService) CreateAgent(ctx context.Context, agent *types.CustomAgent) (*types.CustomAgent, error) {
// Validate required fields
if strings.TrimSpace(agent.Name) == "" {
return nil, ErrAgentNameRequired
}
// Generate UUID and set creation timestamps
if agent.ID == "" {
agent.ID = uuid.New().String()
}
// Get tenant ID from context
tenantID, ok := types.TenantIDFromContext(ctx)
if !ok {
return nil, ErrInvalidTenantID
}
agent.TenantID = tenantID
// Record the creator. Mirrors KnowledgeBase.CreatorID — needed by
// RBAC's RequireOwnershipOrRole so Contributors can edit their own
// agents. Synthetic system-{tenantID} users (X-API-Key path) leave
// the field empty via IsSyntheticUserID, which makes the agent
// tenant-owned (Admin+ only).
if uid, ok := types.UserIDFromContext(ctx); ok && !types.IsSyntheticUserID(uid) {
agent.CreatedBy = uid
}
// Set timestamps
agent.CreatedAt = time.Now()
agent.UpdatedAt = time.Now()
// Ensure agent mode is set for user-created agents
if agent.Config.AgentMode == "" {
agent.Config.AgentMode = types.AgentModeQuickAnswer
}
// Cannot create built-in agents
agent.IsBuiltin = false
// Set defaults
agent.EnsureDefaults()
logger.Infof(ctx, "Creating custom agent, ID: %s, tenant ID: %d, name: %s, agent_mode: %s",
agent.ID, agent.TenantID, agent.Name, agent.Config.AgentMode)
if err := s.repo.CreateAgent(ctx, agent); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agent.ID,
"tenant_id": agent.TenantID,
})
return nil, err
}
logger.Infof(ctx, "Custom agent created successfully, ID: %s, name: %s", agent.ID, agent.Name)
return agent, nil
}
// GetAgentByID retrieves an agent by its ID (including built-in agents)
func (s *customAgentService) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
if id == "" {
logger.Error(ctx, "Agent ID is empty")
return nil, errors.New("agent ID cannot be empty")
}
// Get tenant ID from context
tenantID, ok := types.TenantIDFromContext(ctx)
if !ok {
return nil, ErrInvalidTenantID
}
// Check if it's a built-in agent using the registry
if types.IsBuiltinAgentID(id) {
// Try to get from database first (for customized config)
agent, err := s.repo.GetAgentByID(ctx, id, tenantID)
if err == nil {
// Found in database, return with customized config
return agent, nil
}
// Not in database, return default built-in agent from registry (i18n-aware)
if builtinAgent := types.GetBuiltinAgentWithContext(ctx, id, tenantID); builtinAgent != nil {
return builtinAgent, nil
}
}
// Query from database
agent, err := s.repo.GetAgentByID(ctx, id, tenantID)
if err != nil {
if errors.Is(err, repository.ErrCustomAgentNotFound) {
return nil, ErrAgentNotFound
}
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": id,
})
return nil, err
}
return agent, nil
}
// GetAgentByIDAndTenant retrieves an agent by ID and tenant (for shared agents; does not resolve built-in)
func (s *customAgentService) GetAgentByIDAndTenant(ctx context.Context, id string, tenantID uint64) (*types.CustomAgent, error) {
if id == "" {
logger.Error(ctx, "Agent ID is empty")
return nil, errors.New("agent ID cannot be empty")
}
agent, err := s.repo.GetAgentByID(ctx, id, tenantID)
if err != nil {
if errors.Is(err, repository.ErrCustomAgentNotFound) {
return nil, ErrAgentNotFound
}
return nil, err
}
return agent, nil
}
// ListAgents lists all agents for the current tenant (including built-in agents)
func (s *customAgentService) ListAgents(ctx context.Context) ([]*types.CustomAgent, error) {
tenantID, ok := types.TenantIDFromContext(ctx)
if !ok {
return nil, ErrInvalidTenantID
}
// Get all agents from database (including built-in agents with customized config)
allAgents, err := s.repo.ListAgentsByTenantID(ctx, tenantID)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"tenant_id": tenantID,
})
return nil, err
}
// Track which built-in agents exist in database
builtinInDB := make(map[string]bool)
for _, agent := range allAgents {
if types.IsBuiltinAgentID(agent.ID) {
builtinInDB[agent.ID] = true
}
}
// Build result: built-in agents first, then custom agents
builtinIDs := types.GetBuiltinAgentIDs()
result := make([]*types.CustomAgent, 0, len(allAgents)+len(builtinIDs))
// Add built-in agents in order
for _, builtinID := range builtinIDs {
if builtinInDB[builtinID] {
// Use customized config from database
for _, agent := range allAgents {
if agent.ID == builtinID {
result = append(result, agent)
break
}
}
} else {
// Use default built-in agent (i18n-aware)
if agent := types.GetBuiltinAgentWithContext(ctx, builtinID, tenantID); agent != nil {
result = append(result, agent)
}
}
}
// Add custom agents
for _, agent := range allAgents {
if !types.IsBuiltinAgentID(agent.ID) {
result = append(result, agent)
}
}
return result, nil
}
// UpdateAgent updates an agent's information
func (s *customAgentService) UpdateAgent(ctx context.Context, agent *types.CustomAgent) (*types.CustomAgent, error) {
if agent.ID == "" {
logger.Error(ctx, "Agent ID is empty")
return nil, errors.New("agent ID cannot be empty")
}
// Get tenant ID from context
tenantID, ok := types.TenantIDFromContext(ctx)
if !ok {
return nil, ErrInvalidTenantID
}
// Handle built-in agents specially using registry
if types.IsBuiltinAgentID(agent.ID) {
return s.updateBuiltinAgent(ctx, agent, tenantID)
}
// Get existing agent
existingAgent, err := s.repo.GetAgentByID(ctx, agent.ID, tenantID)
if err != nil {
if errors.Is(err, repository.ErrCustomAgentNotFound) {
return nil, ErrAgentNotFound
}
return nil, err
}
// Cannot modify built-in status
if existingAgent.IsBuiltin {
return nil, ErrCannotModifyBuiltin
}
// Validate name
if strings.TrimSpace(agent.Name) == "" {
return nil, ErrAgentNameRequired
}
// Update fields
existingAgent.Name = agent.Name
existingAgent.Description = agent.Description
existingAgent.Avatar = agent.Avatar
existingAgent.Config = agent.Config
existingAgent.UpdatedAt = time.Now()
// Ensure defaults
existingAgent.EnsureDefaults()
logger.Infof(ctx, "Updating custom agent, ID: %s, name: %s", agent.ID, agent.Name)
if err := s.repo.UpdateAgent(ctx, existingAgent); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agent.ID,
})
return nil, err
}
logger.Infof(ctx, "Custom agent updated successfully, ID: %s", agent.ID)
return existingAgent, nil
}
// updateBuiltinAgent updates a built-in agent's configuration (but not basic info)
func (s *customAgentService) updateBuiltinAgent(ctx context.Context, agent *types.CustomAgent, tenantID uint64) (*types.CustomAgent, error) {
// Get the default built-in agent from registry (i18n-aware)
defaultAgent := types.GetBuiltinAgentWithContext(ctx, agent.ID, tenantID)
if defaultAgent == nil {
return nil, ErrAgentNotFound
}
// Try to get existing customized config from database
existingAgent, err := s.repo.GetAgentByID(ctx, agent.ID, tenantID)
if err != nil && !errors.Is(err, repository.ErrCustomAgentNotFound) {
return nil, err
}
if existingAgent != nil {
// Update existing record - only update config, keep basic info unchanged
existingAgent.Config = agent.Config
existingAgent.UpdatedAt = time.Now()
existingAgent.EnsureDefaults()
logger.Infof(ctx, "Updating built-in agent config, ID: %s", agent.ID)
if err := s.repo.UpdateAgent(ctx, existingAgent); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agent.ID,
})
return nil, err
}
logger.Infof(ctx, "Built-in agent config updated successfully, ID: %s", agent.ID)
return existingAgent, nil
}
// Create new record for built-in agent with customized config
newAgent := &types.CustomAgent{
ID: defaultAgent.ID,
Name: defaultAgent.Name,
Description: defaultAgent.Description,
Avatar: defaultAgent.Avatar,
IsBuiltin: true,
TenantID: tenantID,
Config: agent.Config,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
newAgent.EnsureDefaults()
logger.Infof(ctx, "Creating built-in agent config record, ID: %s, tenant ID: %d", agent.ID, tenantID)
if err := s.repo.CreateAgent(ctx, newAgent); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agent.ID,
"tenant_id": tenantID,
})
return nil, err
}
logger.Infof(ctx, "Built-in agent config record created successfully, ID: %s", agent.ID)
return newAgent, nil
}
// DeleteAgent deletes an agent
func (s *customAgentService) DeleteAgent(ctx context.Context, id string) error {
if id == "" {
logger.Error(ctx, "Agent ID is empty")
return errors.New("agent ID cannot be empty")
}
// Cannot delete built-in agents using registry check
if types.IsBuiltinAgentID(id) {
return ErrCannotDeleteBuiltin
}
// Get tenant ID from context
tenantID, ok := types.TenantIDFromContext(ctx)
if !ok {
return ErrInvalidTenantID
}
// Get existing agent to verify ownership
existingAgent, err := s.repo.GetAgentByID(ctx, id, tenantID)
if err != nil {
if errors.Is(err, repository.ErrCustomAgentNotFound) {
return ErrAgentNotFound
}
return err
}
// Cannot delete built-in agents
if existingAgent.IsBuiltin {
return ErrCannotDeleteBuiltin
}
logger.Infof(ctx, "Deleting custom agent, ID: %s", id)
if err := s.repo.DeleteAgent(ctx, id, tenantID); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": id,
})
return err
}
logger.Infof(ctx, "Custom agent deleted successfully, ID: %s", id)
return nil
}
// CopyAgent creates a copy of an existing agent
func (s *customAgentService) CopyAgent(ctx context.Context, id string) (*types.CustomAgent, error) {
if id == "" {
logger.Error(ctx, "Agent ID is empty")
return nil, errors.New("agent ID cannot be empty")
}
// Get tenant ID from context
tenantID, ok := types.TenantIDFromContext(ctx)
if !ok {
return nil, ErrInvalidTenantID
}
// Get the source agent
sourceAgent, err := s.GetAgentByID(ctx, id)
if err != nil {
return nil, err
}
// Create a new agent with copied data
newAgent := &types.CustomAgent{
ID: uuid.New().String(),
Name: sourceAgent.Name + " (副本)",
Description: sourceAgent.Description,
Avatar: sourceAgent.Avatar,
IsBuiltin: false, // Copied agents are never built-in
TenantID: tenantID,
Config: sourceAgent.Config,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
// The clone is owned by whoever ran the copy, not the original
// creator — same reasoning as CopyKnowledgeBase. Skip synthetic
// API-key users.
if uid, ok := types.UserIDFromContext(ctx); ok && !types.IsSyntheticUserID(uid) {
newAgent.CreatedBy = uid
}
// Ensure defaults
newAgent.EnsureDefaults()
logger.Infof(ctx, "Copying agent, source ID: %s, new ID: %s", id, newAgent.ID)
if err := s.repo.CreateAgent(ctx, newAgent); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"source_agent_id": id,
"new_agent_id": newAgent.ID,
})
return nil, err
}
logger.Infof(ctx, "Agent copied successfully, source ID: %s, new ID: %s", id, newAgent.ID)
return newAgent, nil
}
// GetSuggestedQuestions returns suggested questions for the agent based on its
// associated knowledge bases.
func (s *customAgentService) GetSuggestedQuestions(
ctx context.Context,
agentID string,
kbIDs []string,
knowledgeIDs []string,
limit int,
) ([]types.SuggestedQuestion, error) {
if limit <= 0 {
limit = 6
}
// Get tenant ID from context
tenantID, ok := types.TenantIDFromContext(ctx)
if !ok {
return nil, ErrInvalidTenantID
}
// Get agent configuration
agent, err := s.GetAgentByID(ctx, agentID)
if err != nil {
return nil, err
}
var result []types.SuggestedQuestion
// 1. Add agent config suggested_prompts first (highest priority)
if len(agent.Config.SuggestedPrompts) > 0 {
for _, prompt := range agent.Config.SuggestedPrompts {
if strings.TrimSpace(prompt) == "" {
continue
}
result = append(result, types.SuggestedQuestion{
Question: prompt,
Source: "agent_config",
})
}
}
// 2. Determine knowledge base scope
effectiveKBIDs := kbIDs
if len(effectiveKBIDs) == 0 && len(knowledgeIDs) == 0 {
// Use agent's KB configuration
switch agent.Config.KBSelectionMode {
case "all":
kbs, err := s.kbService.ListKnowledgeBases(ctx)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agentID,
})
// Return what we have so far (agent_config suggestions)
return s.truncateQuestions(result, limit), nil
}
// Honor the agent's implicit/explicit capability requirements so
// e.g. a quick-answer (RAG-only) agent doesn't surface wiki-only
// KBs whose wiki pages it could never answer from. Same filter
// the @ mention dropdown applies on the frontend.
capFilter := tools.DeriveKBFilterForAgent(agent.Config.AgentMode, agent.Config.AllowedTools)
for _, kb := range kbs {
if !capFilter.IsEmpty() &&
!tools.KBSatisfiesAgentRequirements(kb.Capabilities(), agent.Config.AgentMode, agent.Config.AllowedTools) {
continue
}
effectiveKBIDs = append(effectiveKBIDs, kb.ID)
}
case "selected":
effectiveKBIDs = agent.Config.KnowledgeBases
case "none":
// No KB access, return agent_config suggestions only
return s.truncateQuestions(result, limit), nil
default:
// Default to agent's configured KBs
effectiveKBIDs = agent.Config.KnowledgeBases
}
}
if len(effectiveKBIDs) == 0 && len(knowledgeIDs) == 0 {
return s.truncateQuestions(result, limit), nil
}
// Deduplicate questions we've already collected
seen := make(map[string]bool)
for _, q := range result {
seen[q.Question] = true
}
remaining := limit - len(result)
if remaining <= 0 {
return s.truncateQuestions(result, limit), nil
}
// 3. Collect candidate chunks from both FAQ and Document KBs,
// grouped by knowledge_id for diversity.
// knowledgeID -> list of questions
buckets := make(map[string][]types.SuggestedQuestion)
// Determine query scope
queryKBIDs := effectiveKBIDs
queryKnowledgeIDs := knowledgeIDs
// Fetch a large pool so DB-level random sampling covers multiple documents.
fetchLimit := remaining * 5
if fetchLimit < 20 {
fetchLimit = 20
}
// Resolve each KB to the tenant whose chunks should be queried. Cross-tenant
// KBs reached via an organization share map to the source tenant ID; the chunk
// rows live under that tenant. Without this grouping a caller in tenant A
// querying a KB shared from tenant B would hit `tenant_id = A` and get zero
// rows back — the symptom is "suggested questions never appear for shared KBs".
kbGroups := s.groupKBIDsByEffectiveTenant(ctx, tenantID, queryKBIDs)
// Always keep the caller's tenant in the iteration so knowledge_ids-only
// requests (no kbIDs) still execute one query under the caller's tenant.
if len(queryKBIDs) == 0 {
kbGroups[tenantID] = nil
}
// Collect FAQ recommended chunks
for groupTenantID, groupKBIDs := range kbGroups {
faqChunks, err := s.chunkRepo.ListRecommendedFAQChunks(ctx, groupTenantID, groupKBIDs, queryKnowledgeIDs, fetchLimit)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agentID,
"tenant_id": groupTenantID,
})
continue
}
for _, chunk := range faqChunks {
meta, err := chunk.FAQMetadata()
if err != nil || meta == nil || meta.StandardQuestion == "" {
continue
}
if seen[meta.StandardQuestion] {
continue
}
seen[meta.StandardQuestion] = true
buckets[chunk.KnowledgeID] = append(buckets[chunk.KnowledgeID], types.SuggestedQuestion{
Question: meta.StandardQuestion,
Source: "faq",
KnowledgeBaseID: chunk.KnowledgeBaseID,
})
}
}
// Collect Document chunks with generated questions
for groupTenantID, groupKBIDs := range kbGroups {
docChunks, err := s.chunkRepo.ListRecentDocumentChunksWithQuestions(ctx, groupTenantID, groupKBIDs, queryKnowledgeIDs, fetchLimit)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agentID,
"tenant_id": groupTenantID,
})
continue
}
for _, chunk := range docChunks {
meta, err := chunk.DocumentMetadata()
if err != nil || meta == nil || len(meta.GeneratedQuestions) == 0 {
continue
}
q := meta.GeneratedQuestions[0].Question
if q == "" || seen[q] {
continue
}
seen[q] = true
buckets[chunk.KnowledgeID] = append(buckets[chunk.KnowledgeID], types.SuggestedQuestion{
Question: q,
Source: "document",
KnowledgeBaseID: chunk.KnowledgeBaseID,
})
}
}
// Collect Wiki pages as a fallback source. This covers Wiki-only KBs where no
// document chunks carry AI-generated questions (question_generation is skipped
// when the KB does not need an embedding model). knowledge_id filter is
// intentionally ignored here because wiki pages are authored at the KB level
// and are not 1:1 with source knowledge items.
//
// Skip entirely for quick-answer (RAG-only) agents: those can't ever
// retrieve a wiki page, so surfacing wiki-derived suggestions would lure
// the user into asking questions the agent will then answer with empty
// context. Smart-reasoning agents that opt in to wiki tools keep this.
if agent.Config.AgentMode != types.AgentModeQuickAnswer && s.wikiPageRepo != nil {
for groupTenantID, groupKBIDs := range kbGroups {
if len(groupKBIDs) == 0 {
continue
}
wikiPages, err := s.wikiPageRepo.ListRecentForSuggestions(ctx, groupTenantID, groupKBIDs, fetchLimit)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"agent_id": agentID,
"tenant_id": groupTenantID,
})
continue
}
locale, _ := types.LanguageFromContext(ctx)
for _, page := range wikiPages {
q := wikiSuggestionFromPage(page, locale)
if q == "" || seen[q] {
continue
}
seen[q] = true
// Use page.ID as the bucket key so round-robin mixes pages from
// different wiki entries rather than clumping them.
buckets[page.ID] = append(buckets[page.ID], types.SuggestedQuestion{
Question: q,
Source: "wiki",
KnowledgeBaseID: page.KnowledgeBaseID,
})
}
}
}
// 4. Shuffle within each bucket, then round-robin across buckets
// to ensure diversity across different documents.
bucketKeys := make([]string, 0, len(buckets))
for k, qs := range buckets {
bucketKeys = append(bucketKeys, k)
rand.Shuffle(len(qs), func(i, j int) { qs[i], qs[j] = qs[j], qs[i] })
buckets[k] = qs
}
rand.Shuffle(len(bucketKeys), func(i, j int) {
bucketKeys[i], bucketKeys[j] = bucketKeys[j], bucketKeys[i]
})
// Round-robin pick one question from each document in turn.
offsets := make(map[string]int, len(bucketKeys))
for len(result) < limit {
picked := false
for _, key := range bucketKeys {
if len(result) >= limit {
break
}
qs := buckets[key]
idx := offsets[key]
if idx < len(qs) {
result = append(result, qs[idx])
offsets[key] = idx + 1
picked = true
}
}
if !picked {
break
}
}
return s.truncateQuestions(result, limit), nil
}
// truncateQuestions truncates the question list to the specified limit
func (s *customAgentService) truncateQuestions(questions []types.SuggestedQuestion, limit int) []types.SuggestedQuestion {
if len(questions) > limit {
return questions[:limit]
}
return questions
}
// wikiSuggestionFromPage converts a wiki page into a human-readable suggested
// question string. The template is chosen per page type so the chip reads
// naturally for that kind of content:
// - concept: "What is <title>?" works for abstract terms (RAG, embedding,
// idempotency…).
// - entity / summary: "Tell me about <title>" is neutral and works for
// people, places, organizations, products and document summaries where
// "what is <name>?" would read awkwardly ("什么是张三?").
// - everything else (synthesis, comparison, …): the raw title is already a
// good topical query on its own.
func wikiSuggestionFromPage(page *types.WikiPage, locale string) string {
if page == nil {
return ""
}
title := strings.TrimSpace(page.Title)
if title == "" {
return ""
}
switch page.PageType {
case types.WikiPageTypeConcept:
if isEnglishLocale(locale) {
return "What is " + title + "?"
}
return "什么是" + title + "?"
case types.WikiPageTypeEntity, types.WikiPageTypeSummary:
if isEnglishLocale(locale) {
return "Tell me about " + title
}
return "介绍一下" + title
default:
return title
}
}
// groupKBIDsByEffectiveTenant resolves each kbID to the tenant whose chunk
// rows back that KB, so cross-tenant shares can be queried correctly:
// - In-tenant KBs map to the caller's tenant id.
// - KBs owned by another tenant are included only if the caller's tenant
// has at least Viewer access via an organization share, in which case
// the KB maps to its source tenant id (where the chunks actually live).
// - KBs the caller cannot reach (no membership, no share) are silently
// dropped — the suggestion endpoint never returns 403, it just shows
// nothing for that KB, mirroring how search results are scoped.
//
// The result is keyed by effective tenant id so the caller can issue one
// chunk / wiki query per tenant group. Returns an empty (non-nil) map when
// kbIDs is empty.
func (s *customAgentService) groupKBIDsByEffectiveTenant(
ctx context.Context,
callerTenantID uint64,
kbIDs []string,
) map[uint64][]string {
out := make(map[uint64][]string)
if len(kbIDs) == 0 {
return out
}
kbs, err := s.kbService.GetKnowledgeBasesByIDsOnly(ctx, kbIDs)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"kb_ids": kbIDs,
})
// Fall back to caller's tenant so at least in-tenant KBs are queryable;
// chunk repo filtering will drop anything that doesn't match.
out[callerTenantID] = append(out[callerTenantID], kbIDs...)
return out
}
kbByID := make(map[string]*types.KnowledgeBase, len(kbs))
for _, kb := range kbs {
if kb != nil {
kbByID[kb.ID] = kb
}
}
callerRole := types.TenantRoleFromContext(ctx)
for _, kbID := range kbIDs {
kb := kbByID[kbID]
if kb == nil {
continue
}
if kb.TenantID == callerTenantID {
out[callerTenantID] = append(out[callerTenantID], kbID)
continue
}
if s.kbShareService == nil {
continue
}
ok, err := s.kbShareService.HasTenantKBPermission(ctx, kbID, callerTenantID, callerRole, types.OrgRoleViewer)
if err != nil || !ok {
continue
}
out[kb.TenantID] = append(out[kb.TenantID], kbID)
}
return out
}
// isEnglishLocale reports whether the locale string is an English variant.
// Unknown / empty locales fall back to Chinese, matching the product default.
func isEnglishLocale(locale string) bool {
switch locale {
case "en-US", "en", "en-GB":
return true
}
return false
}