Refactor Knowledge Base Collection Initialization and Document Retrieval

- Removed the synchronous preparation of the knowledge base (KB) collection from the InitializeConversation method, now initializing it asynchronously after user login.
- Introduced a new method, GetDocumentsContent, to retrieve content for multiple documents by their IDs, supporting text-based files and improving document handling.
- Updated the API interface to include the new GetDocumentsContent method, enhancing the document management capabilities.
- Enhanced locale handling in the login context to support user preferences during KB collection creation.

agent/assistant/chat.go

@@ -3,126 +3,26 @@ package assistant
 import (
 	"fmt"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/google/uuid"
 	agentcontext "github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/i18n"
 	storetypes "github.com/yaoapp/yao/agent/store/types"
-	"github.com/yaoapp/yao/kb"
-	kbapi "github.com/yaoapp/yao/kb/api"
 )
 
-// kbCollectionCreating tracks collections currently being created to avoid duplicate creation
+// InitializeConversation prepares conversation context (synchronous)
-var kbCollectionCreating sync.Map
+// KB collection is now initialized when user logs in (see openapi/user/login.go)
-
-// InitializeConversation prepares KB collection for the conversation (synchronous)
 func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options ...*agentcontext.Options) error {
-
+	// Reserved for future conversation initialization logic
-	var opts *agentcontext.Options
-	if len(options) > 0 && options[0] != nil {
-		opts = options[0]
-	} else {
-		opts = &agentcontext.Options{}
-	}
-
-	// SKIP: History (for internal calls like title/prompt etc.)
-	if opts.Skip != nil && opts.Skip.History {
-		return nil
-	}
-
-	// Check if authorized info is available
-	if ctx.Authorized == nil {
-		ctx.Logger.Warn("no authorized info, skipping KB collection preparation")
-		return nil
-	}
-
-	// Prepare kb collection
-	err := ast.prepareKBCollection(ctx, opts)
-	if err != nil {
-		// Log but don't fail the chat
-		ctx.Logger.Warn("failed to prepare KB collection: %v", err)
-	}
-
 	return nil
 }
 
-// InitializeConversationAsync prepares KB collection asynchronously
+// InitializeConversationAsync prepares conversation context asynchronously
 func (ast *Assistant) InitializeConversationAsync(ctx *agentcontext.Context, options ...*agentcontext.Options) {
 	go ast.InitializeConversation(ctx, options...)
 }
 
-// prepareKBCollection prepares kb collection (internal method)
-func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agentcontext.Options) error {
-
-	// Get global KB setting
-	kbSetting := GetGlobalKBSetting()
-	if kbSetting == nil || kbSetting.Chat == nil {
-		return nil // No KB configuration for chat, skip
-	}
-
-	// Check if KB API is initialized
-	if kb.API == nil {
-		return fmt.Errorf("KB API not initialized")
-	}
-
-	// Check if authorized info is available
-	if ctx.Authorized == nil {
-		return fmt.Errorf("authorized information not available")
-	}
-
-	chatKB := kbSetting.Chat
-
-	// Debug: log locale information
-	ctx.Logger.Debug("prepareKBCollection: locale=%s", ctx.Locale)
-
-	// Get KB collection ID for this chat session
-	// Same team + user always produces the same ID (idempotent)
-	collectionID := GetChatKBID(ctx.Authorized.TeamID, ctx.Authorized.UserID)
-
-	// Check if this collection is currently being created by another goroutine
-	if _, isCreating := kbCollectionCreating.LoadOrStore(collectionID, true); isCreating {
-		ctx.Logger.Debug("KB collection %s is already being created, skipping", collectionID)
-		return nil
-	}
-	// Ensure cleanup even if panic occurs
-	defer kbCollectionCreating.Delete(collectionID)
-
-	// Check if collection already exists
-	existsResult, err := kb.API.CollectionExists(ctx.Context, collectionID)
-	if err != nil {
-		// If check fails, log and continue to create (let create handle conflicts)
-		ctx.Logger.Warn("failed to check collection existence: %v, will attempt to create", err)
-	} else if existsResult != nil && existsResult.Exists {
-		// Collection exists, no need to create
-		ctx.Logger.Debug("KB collection already exists: %s", collectionID)
-		return nil
-	}
-
-	// Create new collection for this chat session
-	createParams := &kbapi.CreateCollectionParams{
-		ID:                  collectionID,
-		EmbeddingProviderID: chatKB.EmbeddingProviderID,
-		EmbeddingOptionID:   chatKB.EmbeddingOptionID,
-		Locale:              chatKB.Locale,
-		Config:              chatKB.Config,
-		Metadata:            mergeChatMetadata(chatKB.Metadata, ctx),
-		AuthScope:           ctx.Authorized.WithCreateScope(make(map[string]interface{})),
-	}
-
-	_, err = kb.API.CreateCollection(ctx.Context, createParams)
-	if err != nil {
-		return fmt.Errorf("failed to create KB collection: %w", err)
-	}
-
-	ctx.Logger.Info("Created KB collection: %s for team=%s, user=%s",
-		collectionID, ctx.Authorized.TeamID, ctx.Authorized.UserID)
-
-	_ = opts
-	return nil
-}
-
 // GetChatKBID returns the KB collection ID for a chat session
 // Same team + user always returns the same ID (deterministic)
 // Format: chat_{team}_{user} or chat_user_{user} if no team

kb/api/document.go

@@ -3,6 +3,8 @@ package api
 import (
 	"context"
 	"fmt"
+	"os"
+	"strings"
 
 	"github.com/yaoapp/gou/model"
 	"github.com/yaoapp/kun/log"
@@ -277,3 +279,146 @@ func (instance *KBInstance) RemoveDocuments(ctx context.Context, params *RemoveD
 		DBDeletedCount: dbDeletedCount,
 	}, nil
 }
+
+// GetDocumentsContent retrieves content for multiple documents by IDs
+// Returns document info with content (only text-based files are supported)
+func (instance *KBInstance) GetDocumentsContent(ctx context.Context, docIDs []string) ([]map[string]interface{}, error) {
+	if len(docIDs) == 0 {
+		return nil, fmt.Errorf("document IDs are required")
+	}
+
+	// Get document model
+	modelName := "__yao.kb.document"
+	if instance.Config != nil && instance.Config.DocumentModel != "" {
+		modelName = instance.Config.DocumentModel
+	}
+
+	mod := model.Select(modelName)
+	if mod == nil {
+		return nil, fmt.Errorf("document model not found: %s", modelName)
+	}
+
+	results := make([]map[string]interface{}, 0, len(docIDs))
+	for _, docID := range docIDs {
+		param := model.QueryParam{
+			Select: []interface{}{"document_id", "collection_id", "name", "text_content", "type", "status", "file_path", "file_mime_type"},
+			Wheres: []model.QueryWhere{
+				{Column: "document_id", Value: docID},
+			},
+			Limit: 1,
+		}
+
+		docs, err := mod.Get(param)
+		if err != nil {
+			log.Warn("Failed to get document %s: %v", docID, err)
+			continue
+		}
+
+		if len(docs) == 0 {
+			log.Warn("Document not found: %s", docID)
+			continue
+		}
+
+		doc := docs[0]
+		content := ""
+		contentType := "text/plain"
+
+		// Get content type
+		filePath, _ := doc["file_path"].(string)
+		if mimeType, ok := doc["file_mime_type"].(string); ok && mimeType != "" {
+			contentType = mimeType
+		} else if filePath != "" {
+			contentType = inferContentType(filePath)
+		}
+
+		// Only process text-based files
+		if isTextContentType(contentType) {
+			// 1. Try text_content first
+			if textContent, ok := doc["text_content"].(string); ok && textContent != "" {
+				content = textContent
+			} else if filePath != "" {
+				// 2. Read from file_path
+				fileContent, err := readFileContent(filePath)
+				if err != nil {
+					log.Warn("Failed to read file content for %s: %v", docID, err)
+				} else {
+					content = fileContent
+				}
+			}
+		}
+
+		results = append(results, map[string]interface{}{
+			"document_id":   docID,
+			"collection_id": doc["collection_id"],
+			"name":          doc["name"],
+			"content":       content,
+			"content_type":  contentType,
+			"type":          doc["type"],
+			"status":        doc["status"],
+		})
+	}
+
+	return results, nil
+}
+
+// isTextContentType checks if the content type is text-based
+func isTextContentType(contentType string) bool {
+	textTypes := []string{
+		"text/",
+		"application/json",
+		"application/xml",
+		"application/javascript",
+	}
+	for _, tt := range textTypes {
+		if strings.HasPrefix(contentType, tt) {
+			return true
+		}
+	}
+	return false
+}
+
+// inferContentType infers content type from file extension
+func inferContentType(filePath string) string {
+	lower := strings.ToLower(filePath)
+	switch {
+	case strings.HasSuffix(lower, ".md"):
+		return "text/markdown"
+	case strings.HasSuffix(lower, ".txt"):
+		return "text/plain"
+	case strings.HasSuffix(lower, ".html"), strings.HasSuffix(lower, ".htm"):
+		return "text/html"
+	case strings.HasSuffix(lower, ".json"):
+		return "application/json"
+	case strings.HasSuffix(lower, ".xml"):
+		return "application/xml"
+	case strings.HasSuffix(lower, ".csv"):
+		return "text/csv"
+	case strings.HasSuffix(lower, ".pdf"):
+		return "application/pdf"
+	default:
+		return "text/plain"
+	}
+}
+
+// readFileContent reads the content of a file
+func readFileContent(filePath string) (string, error) {
+	// Check if file exists
+	if _, err := os.Stat(filePath); os.IsNotExist(err) {
+		return "", err
+	}
+
+	// Read file content
+	data, err := os.ReadFile(filePath)
+	if err != nil {
+		return "", err
+	}
+
+	// Convert to string and handle encoding
+	content := string(data)
+
+	// Basic cleanup - remove null bytes and normalize line endings
+	content = strings.ReplaceAll(content, "\x00", "")
+	content = strings.ReplaceAll(content, "\r\n", "\n")
+
+	return content, nil
+}

kb/api/interfaces.go

@@ -20,6 +20,7 @@ type API interface {
 	// Document operations
 	ListDocuments(ctx context.Context, filter *ListDocumentsFilter) (*ListDocumentsResult, error)
 	GetDocument(ctx context.Context, docID string, params *GetDocumentParams) (map[string]interface{}, error)
+	GetDocumentsContent(ctx context.Context, docIDs []string) ([]map[string]interface{}, error)
 	RemoveDocuments(ctx context.Context, params *RemoveDocumentsParams) (*RemoveDocumentsResult, error)
 
 	// Document add operations (sync)

openapi/kb/document_process.go

@@ -0,0 +1,70 @@
+package kb
+
+import (
+	"context"
+
+	"github.com/yaoapp/gou/process"
+	"github.com/yaoapp/kun/exception"
+	"github.com/yaoapp/yao/kb"
+)
+
+// ProcessGetDocumentsContent retrieves content for documents by IDs via Yao process
+// Process: kb.documents.getcontents
+//
+// Args[0]: document_ids (string | []string) - Document ID or list of document IDs
+//
+// Returns: []map containing document_id, name, content, content_type, etc. for each document
+//
+// Example:
+//
+//	// Single document
+//	Process("kb.documents.getcontents", "doc_id_123")
+//
+//	// Multiple documents
+//	Process("kb.documents.getcontents", ["doc_id_1", "doc_id_2", "doc_id_3"])
+func ProcessGetDocumentsContent(process *process.Process) interface{} {
+	process.ValidateArgNums(1)
+
+	if kb.API == nil {
+		exception.New("Knowledge base not initialized", 500).Throw()
+	}
+
+	// Support both single string and array of strings
+	var docIDs []string
+	arg := process.Args[0]
+
+	switch v := arg.(type) {
+	case string:
+		if v == "" {
+			exception.New("Document ID is required", 400).Throw()
+		}
+		docIDs = []string{v}
+	case []string:
+		docIDs = v
+	case []interface{}:
+		for _, item := range v {
+			if s, ok := item.(string); ok && s != "" {
+				docIDs = append(docIDs, s)
+			}
+		}
+	default:
+		exception.New("Document IDs must be a string or array of strings", 400).Throw()
+	}
+
+	if len(docIDs) == 0 {
+		exception.New("Document IDs are required", 400).Throw()
+	}
+
+	ctx := process.Context
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	// Call KB API to get documents content
+	results, err := kb.API.GetDocumentsContent(ctx, docIDs)
+	if err != nil {
+		exception.New("Failed to get documents content: "+err.Error(), 500).Throw()
+	}
+
+	return results
+}

openapi/kb/kb.go

@@ -23,6 +23,7 @@ func init() {
 		"documents.addfile":     ProcessAddFile,
 		"documents.addtext":     ProcessAddText,
 		"documents.addurl":      ProcessAddURL,
+		"documents.getcontents": ProcessGetDocumentsContent,
 	})
 }
 

openapi/oauth/types/types.go

@@ -14,6 +14,7 @@ type LoginContext struct {
 	Platform   string `json:"platform,omitempty"`    // Platform (e.g., "ios", "android", "web")
 	Location   string `json:"location,omitempty"`    // Geographic location (optional)
 	RememberMe bool   `json:"remember_me,omitempty"` // Remember Me flag for extended session
+	Locale     string `json:"locale,omitempty"`      // User's preferred locale (e.g., "en-US", "zh-CN")
 }
 
 // MFAOptions contains configuration for MFA operations

openapi/user/login.go

@@ -5,11 +5,15 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/gou/session"
 	"github.com/yaoapp/kun/log"
+	"github.com/yaoapp/yao/agent/assistant"
+	"github.com/yaoapp/yao/kb"
+	kbapi "github.com/yaoapp/yao/kb/api"
 	"github.com/yaoapp/yao/openapi/oauth"
 	"github.com/yaoapp/yao/openapi/oauth/providers/user"
 	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
@@ -18,6 +22,9 @@ import (
 	"github.com/yaoapp/yao/utils/captcha"
 )
 
+// kbCollectionCreating tracks collections currently being created to avoid duplicate creation
+var kbCollectionCreating sync.Map
+
 // getCaptcha is the handler for get captcha image for entry (login/register)
 func getCaptcha(c *gin.Context) {
 	var option captcha.Option = captcha.NewOption()
@@ -268,7 +275,7 @@ func LoginByUserID(userid string, loginCtx *LoginContext) (*LoginResponse, error
 	}
 
 	// Issue tokens without team context
-	return issueTokens(ctx, &IssueTokensParams{
+	resp, err := issueTokens(ctx, &IssueTokensParams{
 		UserID:   userid,
 		TeamID:   "",
 		Team:     nil,
@@ -278,6 +285,18 @@ func LoginByUserID(userid string, loginCtx *LoginContext) (*LoginResponse, error
 		Scopes:   scopes,
 		LoginCtx: loginCtx,
 	})
+	if err != nil {
+		return nil, err
+	}
+
+	// Initialize KB collection asynchronously after successful login
+	locale := ""
+	if loginCtx != nil {
+		locale = loginCtx.Locale
+	}
+	go prepareUserKBCollection(userid, "", locale)
+
+	return resp, nil
 }
 
 // LoginByTeamID is the handler for login by team ID (after team selection)
@@ -311,7 +330,7 @@ func LoginByTeamID(userid string, teamID string, loginCtx *LoginContext) (*Login
 
 	// Handle personal account (no team)
 	if teamID == "" || teamID == "personal" {
-		return issueTokens(ctx, &IssueTokensParams{
+		resp, err := issueTokens(ctx, &IssueTokensParams{
 			UserID:   userid,
 			TeamID:   "",
 			Team:     nil,
@@ -321,6 +340,18 @@ func LoginByTeamID(userid string, teamID string, loginCtx *LoginContext) (*Login
 			Scopes:   scopes,
 			LoginCtx: loginCtx,
 		})
+		if err != nil {
+			return nil, err
+		}
+
+		// Initialize KB collection asynchronously after successful login
+		locale := ""
+		if loginCtx != nil {
+			locale = loginCtx.Locale
+		}
+		go prepareUserKBCollection(userid, "", locale)
+
+		return resp, nil
 	}
 
 	// Verify user is a member of the team and get team details
@@ -346,7 +377,7 @@ func LoginByTeamID(userid string, teamID string, loginCtx *LoginContext) (*Login
 	}
 
 	// Issue tokens with team context and member profile
-	return issueTokens(ctx, &IssueTokensParams{
+	resp, err := issueTokens(ctx, &IssueTokensParams{
 		UserID:   userid,
 		TeamID:   teamID,
 		Team:     team,
@@ -356,6 +387,18 @@ func LoginByTeamID(userid string, teamID string, loginCtx *LoginContext) (*Login
 		Scopes:   scopes,
 		LoginCtx: loginCtx,
 	})
+	if err != nil {
+		return nil, err
+	}
+
+	// Initialize KB collection asynchronously after successful login
+	locale := ""
+	if loginCtx != nil {
+		locale = loginCtx.Locale
+	}
+	go prepareUserKBCollection(userid, teamID, locale)
+
+	return resp, nil
 }
 
 // issueTokens is the core function that issues all necessary tokens (ID token, access token, refresh token)
@@ -602,6 +645,98 @@ func issueTokens(ctx context.Context, params *IssueTokensParams) (*LoginResponse
 	}, nil
 }
 
+// prepareUserKBCollection prepares KB collection for user (called asynchronously after login)
+func prepareUserKBCollection(userID, teamID, locale string) {
+	// Get global KB setting
+	kbSetting := assistant.GetGlobalKBSetting()
+	if kbSetting == nil || kbSetting.Chat == nil {
+		return // No KB configuration for chat, skip
+	}
+
+	// Check if KB API is initialized
+	if kb.API == nil {
+		log.Warn("KB API not initialized, skipping KB collection preparation")
+		return
+	}
+
+	chatKB := kbSetting.Chat
+
+	// Get KB collection ID for this user
+	// Same team + user always produces the same ID (idempotent)
+	collectionID := assistant.GetChatKBID(teamID, userID)
+
+	// Check if this collection is currently being created by another goroutine
+	if _, isCreating := kbCollectionCreating.LoadOrStore(collectionID, true); isCreating {
+		return
+	}
+	// Ensure cleanup even if panic occurs
+	defer kbCollectionCreating.Delete(collectionID)
+
+	// Check if collection already exists
+	ctx := context.Background()
+	existsResult, err := kb.API.CollectionExists(ctx, collectionID)
+	if err != nil {
+		// If check fails, log and continue to create (let create handle conflicts)
+		log.Warn("failed to check collection existence: %v, will attempt to create", err)
+	} else if existsResult != nil && existsResult.Exists {
+		// Collection exists, no need to create
+		return
+	}
+
+	// Build metadata
+	metadata := make(map[string]interface{})
+	for k, v := range chatKB.Metadata {
+		metadata[k] = v
+	}
+	metadata["team_id"] = teamID
+	metadata["user_id"] = userID
+
+	// Ensure name and description are set (required fields)
+	// Use user's locale from login context to determine language
+	isZh := strings.HasPrefix(strings.ToLower(locale), "zh")
+	if _, exists := metadata["name"]; !exists {
+		if isZh {
+			metadata["name"] = "对话知识库"
+		} else {
+			metadata["name"] = "Chat Knowledge Base"
+		}
+	}
+	if _, exists := metadata["description"]; !exists {
+		if isZh {
+			metadata["description"] = "用户对话知识库"
+		} else {
+			metadata["description"] = "User chat knowledge base"
+		}
+	}
+
+	// Build auth scope (use __yao_ prefix for permission fields)
+	authScope := make(map[string]interface{})
+	if teamID != "" {
+		authScope["__yao_team_id"] = teamID
+	}
+	authScope["__yao_created_by"] = userID
+	authScope["__yao_updated_by"] = userID
+
+	// Create new collection for this user
+	createParams := &kbapi.CreateCollectionParams{
+		ID:                  collectionID,
+		EmbeddingProviderID: chatKB.EmbeddingProviderID,
+		EmbeddingOptionID:   chatKB.EmbeddingOptionID,
+		Locale:              chatKB.Locale,
+		Config:              chatKB.Config,
+		Metadata:            metadata,
+		AuthScope:           authScope,
+	}
+
+	_, err = kb.API.CreateCollection(ctx, createParams)
+	if err != nil {
+		log.Warn("failed to create KB collection for user %s: %v", userID, err)
+		return
+	}
+
+	log.Info("Created KB collection: %s for team=%s, user=%s", collectionID, teamID, userID)
+}
+
 // generateSessionID generates a session ID
 func generateSessionID() string {
 	return session.ID()

openapi/user/utils.go

@@ -115,11 +115,22 @@ func makeLoginContext(c *gin.Context) *LoginContext {
 	userAgent := c.GetHeader("User-Agent")
 	device, platform := parseUserAgent(userAgent)
 
+	// Get locale from Accept-Language header or X-Locale header
+	locale := c.GetHeader("X-Locale")
+	if locale == "" {
+		locale = c.GetHeader("Accept-Language")
+		// Parse Accept-Language to get primary language (e.g., "zh-CN,zh;q=0.9" -> "zh-CN")
+		if idx := strings.Index(locale, ","); idx > 0 {
+			locale = locale[:idx]
+		}
+	}
+
 	return &LoginContext{
 		IP:        userIPAddress(c),
 		UserAgent: userAgent,
 		Device:    device,
 		Platform:  platform,
+		Locale:    locale,
 	}
 }