diff --git a/kb/kb.go b/kb/kb.go index a29ddd5d..f4c75fbc 100644 --- a/kb/kb.go +++ b/kb/kb.go @@ -150,3 +150,17 @@ func GetProviderWithLanguage(typ string, id string, locale string) (*kbtypes.Pro return knowledgeBase.Providers.GetProvider(typ, id, locale) } + +// GetConfig returns the knowledge base configuration +func GetConfig() (*kbtypes.Config, error) { + if Instance == nil { + return nil, fmt.Errorf("knowledge base not initialized") + } + + knowledgeBase, ok := Instance.(*KnowledgeBase) + if !ok { + return nil, fmt.Errorf("knowledge base not initialized") + } + + return knowledgeBase.Config, nil +} diff --git a/kb/types/collection.go b/kb/types/collection.go new file mode 100644 index 00000000..26e19d5f --- /dev/null +++ b/kb/types/collection.go @@ -0,0 +1,110 @@ +package types + +import ( + "fmt" + + "github.com/yaoapp/gou/model" + "github.com/yaoapp/kun/maps" +) + +// SearchCollections searches collections with pagination +func (c *Config) SearchCollections(param model.QueryParam, page int, pagesize int) (maps.MapStr, error) { + modelName := c.CollectionModel + if modelName == "" { + modelName = "__yao.kb.collection" + } + + mod := model.Select(modelName) + if mod == nil { + return nil, fmt.Errorf("collection model not found: %s", modelName) + } + return mod.Paginate(param, page, pagesize) +} + +// FindCollection finds a single collection by collection_id +func (c *Config) FindCollection(collectionID string, param model.QueryParam) (maps.MapStr, error) { + modelName := c.CollectionModel + if modelName == "" { + modelName = "__yao.kb.collection" + } + + mod := model.Select(modelName) + if mod == nil { + return nil, fmt.Errorf("collection model not found: %s", modelName) + } + + param.Wheres = append(param.Wheres, model.QueryWhere{ + Column: "collection_id", + Value: collectionID, + }) + param.Limit = 1 + + res, err := mod.Get(param) + if err != nil { + return nil, err + } + if len(res) == 0 { + return nil, fmt.Errorf("collection not found: %s", collectionID) + } + return res[0], nil +} + +// CreateCollection creates a new collection record +func (c *Config) CreateCollection(data maps.MapStrAny) (int, error) { + modelName := c.CollectionModel + if modelName == "" { + modelName = "__yao.kb.collection" + } + + mod := model.Select(modelName) + if mod == nil { + return 0, fmt.Errorf("collection model not found: %s", modelName) + } + return mod.Create(data) +} + +// UpdateCollection updates a collection by collection_id +func (c *Config) UpdateCollection(collectionID string, data maps.MapStrAny) error { + modelName := c.CollectionModel + if modelName == "" { + modelName = "__yao.kb.collection" + } + + mod := model.Select(modelName) + if mod == nil { + return fmt.Errorf("collection model not found: %s", modelName) + } + + param := model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "collection_id", Value: collectionID}, + }, + Limit: 1, + } + + _, err := mod.UpdateWhere(param, data) + return err +} + +// RemoveCollection removes a collection by collection_id +func (c *Config) RemoveCollection(collectionID string) error { + modelName := c.CollectionModel + if modelName == "" { + modelName = "__yao.kb.collection" + } + + mod := model.Select(modelName) + if mod == nil { + return fmt.Errorf("collection model not found: %s", modelName) + } + + param := model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "collection_id", Value: collectionID}, + }, + Limit: 1, + } + + _, err := mod.DeleteWhere(param) + return err +} diff --git a/kb/types/config.go b/kb/types/config.go index 97a4dcf8..c27fd92d 100644 --- a/kb/types/config.go +++ b/kb/types/config.go @@ -323,6 +323,16 @@ func (c *Config) UnmarshalJSON(data []byte) error { c.Uploader = "__yao.attachment" } + // Set default collection model if not configured + if c.CollectionModel == "" { + c.CollectionModel = "__yao.kb.collection" + } + + // Set default document model if not configured + if c.DocumentModel == "" { + c.DocumentModel = "__yao.kb.document" + } + // Note: Features will be computed later after providers are loaded return nil } diff --git a/kb/types/document.go b/kb/types/document.go new file mode 100644 index 00000000..cbe8faec --- /dev/null +++ b/kb/types/document.go @@ -0,0 +1,110 @@ +package types + +import ( + "fmt" + + "github.com/yaoapp/gou/model" + "github.com/yaoapp/kun/maps" +) + +// SearchDocuments searches documents with pagination +func (c *Config) SearchDocuments(param model.QueryParam, page int, pagesize int) (maps.MapStr, error) { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return nil, fmt.Errorf("document model not found: %s", modelName) + } + return mod.Paginate(param, page, pagesize) +} + +// FindDocument finds a single document by document_id +func (c *Config) FindDocument(documentID string, param model.QueryParam) (maps.MapStr, error) { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return nil, fmt.Errorf("document model not found: %s", modelName) + } + + param.Wheres = append(param.Wheres, model.QueryWhere{ + Column: "document_id", + Value: documentID, + }) + param.Limit = 1 + + res, err := mod.Get(param) + if err != nil { + return nil, err + } + if len(res) == 0 { + return nil, fmt.Errorf("document not found: %s", documentID) + } + return res[0], nil +} + +// CreateDocument creates a new document record +func (c *Config) CreateDocument(data maps.MapStrAny) (int, error) { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return 0, fmt.Errorf("document model not found: %s", modelName) + } + return mod.Create(data) +} + +// UpdateDocument updates a document by document_id +func (c *Config) UpdateDocument(documentID string, data maps.MapStrAny) error { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return fmt.Errorf("document model not found: %s", modelName) + } + + param := model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "document_id", Value: documentID}, + }, + Limit: 1, + } + + _, err := mod.UpdateWhere(param, data) + return err +} + +// RemoveDocument removes a document by document_id +func (c *Config) RemoveDocument(documentID string) error { + modelName := c.DocumentModel + if modelName == "" { + modelName = "__yao.kb.document" + } + + mod := model.Select(modelName) + if mod == nil { + return fmt.Errorf("document model not found: %s", modelName) + } + + param := model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "document_id", Value: documentID}, + }, + Limit: 1, + } + + _, err := mod.DeleteWhere(param) + return err +} diff --git a/kb/types/types.go b/kb/types/types.go index 4011d2b4..6fe2e887 100644 --- a/kb/types/types.go +++ b/kb/types/types.go @@ -63,6 +63,12 @@ type Config struct { // KV store name (Optional with default value) Store string `json:"store,omitempty" yaml:"store,omitempty"` // Default: "__yao.kb.store" + // Bind Collection Model + CollectionModel string `json:"collection_model,omitempty" yaml:"collection_model,omitempty"` // Default: "__yao.kb.collection" + + // Bind Document Model + DocumentModel string `json:"document_model,omitempty" yaml:"document_model,omitempty"` // Default: "__yao.kb.document" + // PDF parser configuration (Optional) PDF *PDFConfig `json:"pdf,omitempty" yaml:"pdf,omitempty"` diff --git a/openapi/kb/collection.go b/openapi/kb/collection.go index c9bf973c..f546aa54 100644 --- a/openapi/kb/collection.go +++ b/openapi/kb/collection.go @@ -5,6 +5,8 @@ import ( "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/kun/maps" "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/openapi/response" ) @@ -20,44 +22,9 @@ type ProviderSettings struct { // CreateCollection creates a new collection func CreateCollection(c *gin.Context) { - var req CreateCollectionRequest - - // Parse and bind JSON request - if err := c.ShouldBindJSON(&req); err != nil { - // Create a custom error with the same structure but specific message - errorResp := &response.ErrorResponse{ - Code: response.ErrInvalidRequest.Code, - ErrorDescription: "Invalid request format: " + err.Error(), - } - response.RespondWithError(c, response.StatusBadRequest, errorResp) - return - } - - // Get provider settings by provider id and option value - providerSettings, err := getProviderSettings(req.Config.EmbeddingProvider, req.Config.EmbeddingOption, req.Config.Locale) + // Prepare request and database data + req, collectionData, err := PrepareCreateCollection(c) if err != nil { - errorResp := &response.ErrorResponse{ - Code: response.ErrInvalidRequest.Code, - ErrorDescription: fmt.Sprintf("Failed to resolve provider settings: %v", err), - } - response.RespondWithError(c, response.StatusBadRequest, errorResp) - return - } - - // Set dimension by provider settings and add original provider id and option value to metadata with prefix __ - req.Config.Dimension = providerSettings.Dimension - if req.Metadata == nil { - req.Metadata = make(map[string]interface{}) - } - req.Metadata["__embedding_provider"] = req.Config.EmbeddingProvider - req.Metadata["__embedding_option"] = req.Config.EmbeddingOption - if req.Config.Locale != "" { - req.Metadata["__locale"] = req.Config.Locale - } - - // Validate request parameters - if err := validateCreateCollectionRequest(&req); err != nil { - // Create a custom error with the same structure but specific message errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: err.Error(), @@ -68,7 +35,6 @@ func CreateCollection(c *gin.Context) { // Check if kb.Instance is available if kb.Instance == nil { - // Create a custom error with the same structure but specific message errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Knowledge base not initialized", @@ -77,7 +43,29 @@ func CreateCollection(c *gin.Context) { return } - // Create CollectionConfig + // Get KB config + config, err := kb.GetConfig() + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to get KB config: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // First create database record + _, err = config.CreateCollection(maps.MapStrAny(collectionData)) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to save collection metadata: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // Create CollectionConfig for GraphRag collectionConfig := types.CollectionConfig{ ID: req.ID, Metadata: req.Metadata, @@ -87,7 +75,12 @@ func CreateCollection(c *gin.Context) { // Call the actual CreateCollection method collectionID, err := kb.Instance.CreateCollection(c.Request.Context(), collectionConfig) if err != nil { - // Create a custom error with the same structure but specific message + // Rollback: remove the database record + rollbackErr := config.RemoveCollection(req.ID) + if rollbackErr != nil { + log.Error("Failed to rollback collection database record: %v", rollbackErr) + } + errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to create collection: " + err.Error(), @@ -96,6 +89,12 @@ func CreateCollection(c *gin.Context) { return } + // Update status to active after successful creation + updateErr := config.UpdateCollection(req.ID, maps.MapStrAny{"status": "active"}) + if updateErr != nil { + log.Error("Failed to update collection status to active: %v", updateErr) + } + successData := gin.H{ "message": "Collection created successfully", "collection_id": collectionID, @@ -146,6 +145,13 @@ func RemoveCollection(c *gin.Context) { return } + // Remove collection from database after successful GraphRag removal + if config, err := kb.GetConfig(); err == nil { + if err := config.RemoveCollection(collectionID); err != nil { + log.Error("Failed to remove collection from database: %v", err) + } + } + successData := gin.H{ "message": "Collection removed successfully", "collection_id": collectionID, @@ -289,6 +295,27 @@ func UpdateCollectionMetadata(c *gin.Context) { return } + // Update collection metadata in database after successful GraphRag update + if config, err := kb.GetConfig(); err == nil { + // Prepare update data from metadata + updateData := maps.MapStrAny{} + if name, ok := req.Metadata["name"]; ok { + updateData["name"] = name + } + if description, ok := req.Metadata["description"]; ok { + updateData["description"] = description + } + if status, ok := req.Metadata["status"]; ok { + updateData["status"] = status + } + + if len(updateData) > 0 { + if err := config.UpdateCollection(collectionID, updateData); err != nil { + log.Error("Failed to update collection in database: %v", err) + } + } + } + successData := gin.H{ "message": "Collection metadata updated successfully", "collection_id": collectionID, diff --git a/openapi/kb/document.go b/openapi/kb/document.go index df463d7f..66a3dc0b 100644 --- a/openapi/kb/document.go +++ b/openapi/kb/document.go @@ -6,6 +6,8 @@ import ( "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/kun/maps" "github.com/yaoapp/yao/attachment" "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/openapi/response" @@ -120,49 +122,88 @@ func handleAsync(c *gin.Context, syncHandler func(*gin.Context)) { // AddFile adds a file to a collection func AddFile(c *gin.Context) { - var req AddFileRequest - // Check if kb.Instance is available if !checkKBInstance(c) { return } - // Validate request - if err := validateRequest(c, &req); err != nil { + // Prepare request and database data + req, documentData, err := PrepareAddFile(c) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) return } - // Validate file and get path - path, contentType, err := validateFileAndGetPath(c, &req) + // Get KB config + config, err := kb.GetConfig() if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to get KB config: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // First create database record + _, err = config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to save document metadata: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Convert request to UpsertOptions + path, contentType, err := validateFileAndGetPath(c, req) + if err != nil { + // Rollback: remove the database record + if err := config.RemoveDocument(req.DocID); err != nil { + log.Error("Failed to rollback document database record: %v", err) + } + return + } + upsertOptions, err := getUpsertOptions(c, &req.BaseUpsertRequest, path, contentType) if err != nil { + // Rollback: remove the database record + if err := config.RemoveDocument(req.DocID); err != nil { + log.Error("Failed to rollback document database record: %v", err) + } return } // Perform upsert operation with file ID - // Note: In a real implementation, you would need to fetch the file content - // using req.FileID and pass it to the upsert operation - docID, err := kb.Instance.AddFile(c.Request.Context(), req.FileID, upsertOptions) + _, err = kb.Instance.AddFile(c.Request.Context(), req.FileID, upsertOptions) if err != nil { + // Update status to error and return error response + config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to upsert file: " + err.Error(), + ErrorDescription: "Failed to add file: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } + // Update status to completed after successful processing + if err := config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "completed"}); err != nil { + log.Error("Failed to update document status to completed: %v", err) + } + // Return success response result := gin.H{ "message": "File added successfully", "collection_id": req.CollectionID, "file_id": req.FileID, - "doc_id": docID, + "doc_id": req.DocID, } response.RespondWithSuccess(c, response.StatusCreated, result) @@ -200,40 +241,78 @@ func AddFileAsync(c *gin.Context) { // AddText adds text to a collection func AddText(c *gin.Context) { - var req AddTextRequest - - // Validate request - if err := validateRequest(c, &req); err != nil { + // Check if kb.Instance is available + if !checkKBInstance(c) { return } - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Prepare request and database data + req, documentData, err := PrepareAddText(c) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Get KB config + config, err := kb.GetConfig() + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to get KB config: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // First create database record + _, err = config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to save document metadata: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Convert request to UpsertOptions upsertOptions, err := getUpsertOptions(c, &req.BaseUpsertRequest) if err != nil { + // Rollback: remove the database record + if err := config.RemoveDocument(req.DocID); err != nil { + log.Error("Failed to rollback document database record: %v", err) + } return } // Perform upsert operation with text - docID, err := kb.Instance.AddText(c.Request.Context(), req.Text, upsertOptions) + _, err = kb.Instance.AddText(c.Request.Context(), req.Text, upsertOptions) if err != nil { + // Update status to error and return error response + config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to upsert text: " + err.Error(), + ErrorDescription: "Failed to add text: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } + // Update status to completed after successful processing + if err := config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "completed"}); err != nil { + log.Error("Failed to update document status to completed: %v", err) + } + // Return success response result := gin.H{ "message": "Text added successfully", "collection_id": req.CollectionID, - "doc_id": docID, + "doc_id": req.DocID, } response.RespondWithSuccess(c, response.StatusCreated, result) @@ -265,41 +344,79 @@ func AddTextAsync(c *gin.Context) { // AddURL adds a URL to a collection func AddURL(c *gin.Context) { - var req AddURLRequest - - // Validate request - if err := validateRequest(c, &req); err != nil { + // Check if kb.Instance is available + if !checkKBInstance(c) { return } - // Check if kb.Instance is available - if !checkKBInstance(c) { + // Prepare request and database data + req, documentData, err := PrepareAddURL(c) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrInvalidRequest.Code, + ErrorDescription: err.Error(), + } + response.RespondWithError(c, response.StatusBadRequest, errorResp) + return + } + + // Get KB config + config, err := kb.GetConfig() + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to get KB config: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) + return + } + + // First create database record + _, err = config.CreateDocument(maps.MapStrAny(documentData)) + if err != nil { + errorResp := &response.ErrorResponse{ + Code: response.ErrServerError.Code, + ErrorDescription: "Failed to save document metadata: " + err.Error(), + } + response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Convert request to UpsertOptions upsertOptions, err := getUpsertOptions(c, &req.BaseUpsertRequest) if err != nil { + // Rollback: remove the database record + if err := config.RemoveDocument(req.DocID); err != nil { + log.Error("Failed to rollback document database record: %v", err) + } return } // Perform upsert operation with URL - docID, err := kb.Instance.AddURL(c.Request.Context(), req.URL, upsertOptions) + _, err = kb.Instance.AddURL(c.Request.Context(), req.URL, upsertOptions) if err != nil { + // Update status to error and return error response + config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "error", "error_message": err.Error()}) + errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to upsert URL: " + err.Error(), + ErrorDescription: "Failed to add URL: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } + // Update status to completed after successful processing + if err := config.UpdateDocument(req.DocID, maps.MapStrAny{"status": "completed"}); err != nil { + log.Error("Failed to update document status to completed: %v", err) + } + // Return success response result := gin.H{ "message": "URL added successfully", "collection_id": req.CollectionID, "url": req.URL, - "doc_id": docID, + "doc_id": req.DocID, } response.RespondWithSuccess(c, response.StatusCreated, result) diff --git a/openapi/kb/utils.go b/openapi/kb/utils.go index 0bd10531..7ba2ac0c 100644 --- a/openapi/kb/utils.go +++ b/openapi/kb/utils.go @@ -3,7 +3,10 @@ package kb import ( "fmt" + "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/types" + "github.com/yaoapp/gou/graphrag/utils" + "github.com/yaoapp/yao/attachment" "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/kb/providers/factory" kbtypes "github.com/yaoapp/yao/kb/types" @@ -399,3 +402,241 @@ func (r *UpdateSegmentsRequest) Validate() error { } return nil } + +// PrepareCreateCollection prepares CreateCollection request and database data +func PrepareCreateCollection(c *gin.Context) (*CreateCollectionRequest, map[string]interface{}, error) { + var req CreateCollectionRequest + + // Parse and bind JSON request + if err := c.ShouldBindJSON(&req); err != nil { + return nil, nil, fmt.Errorf("invalid request format: %w", err) + } + + // Get provider settings first to resolve dimension + providerSettings, err := getProviderSettings(req.Config.EmbeddingProvider, req.Config.EmbeddingOption, req.Config.Locale) + if err != nil { + return nil, nil, fmt.Errorf("failed to resolve provider settings: %w", err) + } + + // Set dimension from provider settings + req.Config.Dimension = providerSettings.Dimension + + // Add metadata with provider information + if req.Metadata == nil { + req.Metadata = make(map[string]interface{}) + } + req.Metadata["__embedding_provider"] = req.Config.EmbeddingProvider + req.Metadata["__embedding_option"] = req.Config.EmbeddingOption + if req.Config.Locale != "" { + req.Metadata["__locale"] = req.Config.Locale + } + + // Now validate request parameters (after dimension and metadata are set) + if err := validateCreateCollectionRequest(&req); err != nil { + return nil, nil, err + } + + // Prepare collection data for database + data := map[string]interface{}{ + "collection_id": req.ID, + "name": req.Metadata["name"], + "description": req.Metadata["description"], + "status": "creating", + "embedding_provider": req.Config.EmbeddingProvider, + "embedding_option": req.Config.EmbeddingOption, + "locale": req.Config.Locale, + "distance": req.Config.Distance, + "index_type": req.Config.IndexType, + } + + // Add optional HNSW parameters + if req.Config.M > 0 { + data["m"] = req.Config.M + } + if req.Config.EfConstruction > 0 { + data["ef_construction"] = req.Config.EfConstruction + } + if req.Config.EfSearch > 0 { + data["ef_search"] = req.Config.EfSearch + } + + // Add optional IVF parameters + if req.Config.NumLists > 0 { + data["num_lists"] = req.Config.NumLists + } + if req.Config.NumProbes > 0 { + data["num_probes"] = req.Config.NumProbes + } + + // Add context fields (permissions, user info, etc.) + addContextFields(c, data) + + return &req, data, nil +} + +// PrepareAddFile prepares AddFile request and database data +func PrepareAddFile(c *gin.Context) (*AddFileRequest, map[string]interface{}, error) { + var req AddFileRequest + + // Parse and validate request + if err := validateRequest(c, &req); err != nil { + return nil, nil, err + } + + // Validate file and get path + path, contentType, err := validateFileAndGetPath(c, &req) + if err != nil { + return nil, nil, err + } + + // Get file info + m, _ := attachment.Managers[req.Uploader] + fileInfo, _ := m.Info(c.Request.Context(), req.FileID) + + // Generate document ID if not provided + if req.DocID == "" { + req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) + } + + // Prepare document data for database + data := map[string]interface{}{ + "document_id": req.DocID, + "collection_id": req.CollectionID, + "name": fileInfo.Filename, + "type": "file", + "status": "pending", + "uploader_id": req.Uploader, + "file_name": fileInfo.Filename, + "file_path": path, + "file_mime_type": contentType, + "size": int64(fileInfo.Bytes), + } + + addBaseRequestFields(data, &req.BaseUpsertRequest) + addContextFields(c, data) + + return &req, data, nil +} + +// PrepareAddText prepares AddText request and database data +func PrepareAddText(c *gin.Context) (*AddTextRequest, map[string]interface{}, error) { + var req AddTextRequest + + // Parse and validate request + if err := validateRequest(c, &req); err != nil { + return nil, nil, err + } + + // Generate document ID if not provided + if req.DocID == "" { + req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) + } + + // Prepare document data for database + data := map[string]interface{}{ + "document_id": req.DocID, + "collection_id": req.CollectionID, + "name": "Text Document", + "type": "text", + "status": "pending", + "text_content": req.Text, + "size": int64(len(req.Text)), + } + + // Use title from metadata if available + if req.Metadata != nil { + if title, ok := req.Metadata["title"].(string); ok && title != "" { + data["name"] = title + } + } + + addBaseRequestFields(data, &req.BaseUpsertRequest) + addContextFields(c, data) + + return &req, data, nil +} + +// PrepareAddURL prepares AddURL request and database data +func PrepareAddURL(c *gin.Context) (*AddURLRequest, map[string]interface{}, error) { + var req AddURLRequest + + // Parse and validate request + if err := validateRequest(c, &req); err != nil { + return nil, nil, err + } + + // Generate document ID if not provided + if req.DocID == "" { + req.DocID = utils.GenDocIDWithCollectionID(req.CollectionID) + } + + // Prepare document data for database + data := map[string]interface{}{ + "document_id": req.DocID, + "collection_id": req.CollectionID, + "name": req.URL, + "type": "url", + "status": "pending", + "url": req.URL, + } + + // Use title from metadata if available + if req.Metadata != nil { + if title, ok := req.Metadata["title"].(string); ok && title != "" { + data["name"] = title + data["url_title"] = title + } + } + + addBaseRequestFields(data, &req.BaseUpsertRequest) + addContextFields(c, data) + + return &req, data, nil +} + +// addBaseRequestFields adds common fields from BaseUpsertRequest +func addBaseRequestFields(data map[string]interface{}, req *BaseUpsertRequest) { + if req.Locale != "" { + data["locale"] = req.Locale + } + if req.DocID != "" { + data["document_id"] = req.DocID + } + if req.Metadata != nil { + data["tags"] = req.Metadata + } + + // Add provider configurations + if req.Converter != nil { + data["converter_provider_id"] = req.Converter.ProviderID + if req.Converter.Option != nil { + data["converter_properties"] = req.Converter.Option.Properties + } + } + if req.Fetcher != nil { + data["fetcher_provider_id"] = req.Fetcher.ProviderID + if req.Fetcher.Option != nil { + data["fetcher_properties"] = req.Fetcher.Option.Properties + } + } + if req.Chunking != nil { + data["chunking_provider_id"] = req.Chunking.ProviderID + if req.Chunking.Option != nil { + data["chunking_properties"] = req.Chunking.Option.Properties + } + } + if req.Extraction != nil { + data["extractor_provider_id"] = req.Extraction.ProviderID + if req.Extraction.Option != nil { + data["extractor_properties"] = req.Extraction.Option.Properties + } + } +} + +// addContextFields adds context-specific fields like permissions, user info +func addContextFields(c *gin.Context, data map[string]interface{}) { + // TODO: Add permission-related fields from Guard + // Example: data["user_id"] = c.GetString("user_id") + // Example: data["permissions"] = c.Get("permissions") + // Example: data["tenant_id"] = c.GetString("tenant_id") +}