diff --git a/agent/assistant/chat.go b/agent/assistant/chat.go index f714950c..9feceab7 100644 --- a/agent/assistant/chat.go +++ b/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 -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 { - 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) - } - + // Reserved for future conversation initialization logic 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 diff --git a/kb/api/document.go b/kb/api/document.go index faa6d3aa..aad111c0 100644 --- a/kb/api/document.go +++ b/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 +} diff --git a/kb/api/interfaces.go b/kb/api/interfaces.go index 3cab0013..ad263c7f 100644 --- a/kb/api/interfaces.go +++ b/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) diff --git a/openapi/kb/document_process.go b/openapi/kb/document_process.go new file mode 100644 index 00000000..a334006d --- /dev/null +++ b/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 +} diff --git a/openapi/kb/kb.go b/openapi/kb/kb.go index 7f1cb5ca..baa474d6 100644 --- a/openapi/kb/kb.go +++ b/openapi/kb/kb.go @@ -20,9 +20,10 @@ func init() { "collection.updatemetadata": ProcessUpdateCollectionMetadata, // Document processes - "documents.addfile": ProcessAddFile, - "documents.addtext": ProcessAddText, - "documents.addurl": ProcessAddURL, + "documents.addfile": ProcessAddFile, + "documents.addtext": ProcessAddText, + "documents.addurl": ProcessAddURL, + "documents.getcontents": ProcessGetDocumentsContent, }) } diff --git a/openapi/oauth/types/types.go b/openapi/oauth/types/types.go index 8bd573a3..730cee42 100644 --- a/openapi/oauth/types/types.go +++ b/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 diff --git a/openapi/user/login.go b/openapi/user/login.go index 2f156e95..a7c4dc4e 100644 --- a/openapi/user/login.go +++ b/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() diff --git a/openapi/user/utils.go b/openapi/user/utils.go index 3fe448c4..5f97f40b 100644 --- a/openapi/user/utils.go +++ b/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, } }