Merge pull request #1101 from trheyi/main

Implement knowledge base configuration retrieval

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

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

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
 }

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

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"`
 

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,

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)

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")
+}