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.
This commit is contained in:
Max 2026-01-10 12:38:42 +08:00
parent ff667b69d9
commit 4ba62600dc
8 changed files with 374 additions and 110 deletions

View file

@ -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
var kbCollectionCreating sync.Map
// InitializeConversation prepares KB collection for the conversation (synchronous)
// InitializeConversation prepares conversation context (synchronous)
// KB collection is now initialized when user logs in (see openapi/user/login.go)
func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options ...*agentcontext.Options) error {
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 {
// Reserved for future conversation initialization logic
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

View file

@ -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
}

View file

@ -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)

View file

@ -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
}

View file

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

View file

@ -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

View file

@ -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()

View file

@ -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,
}
}