Implement knowledge base configuration retrieval and enhance collection/document handling
- Added GetConfig function to retrieve the knowledge base configuration, ensuring proper initialization checks. - Refactored CreateCollection, AddFile, AddText, and AddURL functions to utilize the new GetConfig method for improved error handling and database record management. - Introduced preparation functions (PrepareCreateCollection, PrepareAddFile, PrepareAddText, PrepareAddURL) to streamline request handling and database data preparation. - Enhanced error responses and rollback mechanisms for document and collection operations, improving robustness and clarity in error handling.
This commit is contained in:
Max
parent
d3e649ebfe
commit
71b78bc23b
6 changed files with 688 additions and 69 deletions
14
kb/kb.go
14
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
|
||||
}
|
||||
|
|
|
|||
110
kb/types/collection.go
Normal file
110
kb/types/collection.go
Normal file
|
|
@ -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
|
||||
}
|
||||
110
kb/types/document.go
Normal file
110
kb/types/document.go
Normal file
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue