Implement document management handlers for listing, retrieving, and removing documents
- Added ListDocuments and ScrollDocuments functions for paginated document retrieval. - Implemented GetDocument function to fetch document details by ID, including error handling for missing IDs. - Introduced RemoveDocs function to handle document deletion requests. - Removed outdated AddFile, AddText, and AddURL functions to streamline document management logic.
This commit is contained in:
Max
parent
6acea5b004
commit
40510ec208
6 changed files with 842 additions and 811 deletions
128
openapi/kb/addfile.go
Normal file
128
openapi/kb/addfile.go
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
package kb
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/kun/log"
|
||||
"github.com/yaoapp/kun/maps"
|
||||
"github.com/yaoapp/yao/kb"
|
||||
"github.com/yaoapp/yao/openapi/response"
|
||||
)
|
||||
|
||||
// AddFile adds a file to a collection
|
||||
func AddFile(c *gin.Context) {
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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
|
||||
_, 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 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": req.DocID,
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusCreated, result)
|
||||
}
|
||||
|
||||
// AddFileAsync adds file to a collection asynchronously
|
||||
func AddFileAsync(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 {
|
||||
return
|
||||
}
|
||||
|
||||
// Validate file and get path
|
||||
_, _, err := validateFileAndGetPath(c, &req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert request to UpsertOptions (just for validation)
|
||||
_, err = getUpsertOptions(c, &req.BaseUpsertRequest)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle async processing
|
||||
handleAsync(c, AddFile)
|
||||
}
|
||||
112
openapi/kb/addtext.go
Normal file
112
openapi/kb/addtext.go
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
package kb
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/kun/log"
|
||||
"github.com/yaoapp/kun/maps"
|
||||
"github.com/yaoapp/yao/kb"
|
||||
"github.com/yaoapp/yao/openapi/response"
|
||||
)
|
||||
|
||||
// AddText adds text to a collection
|
||||
func AddText(c *gin.Context) {
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// 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
|
||||
_, 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 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": req.DocID,
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusCreated, result)
|
||||
}
|
||||
|
||||
// AddTextAsync adds text to a collection asynchronously
|
||||
func AddTextAsync(c *gin.Context) {
|
||||
var req AddTextRequest
|
||||
|
||||
// Validate request
|
||||
if err := validateRequest(c, &req); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert request to UpsertOptions (just for validation)
|
||||
_, err := getUpsertOptions(c, &req.BaseUpsertRequest)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle async processing
|
||||
handleAsync(c, AddText)
|
||||
}
|
||||
113
openapi/kb/addurl.go
Normal file
113
openapi/kb/addurl.go
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
package kb
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/kun/log"
|
||||
"github.com/yaoapp/kun/maps"
|
||||
"github.com/yaoapp/yao/kb"
|
||||
"github.com/yaoapp/yao/openapi/response"
|
||||
)
|
||||
|
||||
// AddURL adds a URL to a collection
|
||||
func AddURL(c *gin.Context) {
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// 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
|
||||
_, 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 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": req.DocID,
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusCreated, result)
|
||||
}
|
||||
|
||||
// AddURLAsync adds a URL to a collection asynchronously
|
||||
func AddURLAsync(c *gin.Context) {
|
||||
var req AddURLRequest
|
||||
|
||||
// Validate request
|
||||
if err := validateRequest(c, &req); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert request to UpsertOptions (just for validation)
|
||||
_, err := getUpsertOptions(c, &req.BaseUpsertRequest)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle async processing
|
||||
handleAsync(c, AddURL)
|
||||
}
|
||||
|
|
@ -6,8 +6,6 @@ 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"
|
||||
|
|
@ -15,6 +13,55 @@ import (
|
|||
|
||||
// Document Management Handlers
|
||||
|
||||
// ListDocuments lists documents with pagination
|
||||
func ListDocuments(c *gin.Context) {
|
||||
// TODO: Implement list documents logic
|
||||
// Query parameters for pagination: page, limit, filter, etc.
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"documents": []interface{}{},
|
||||
"total": 0,
|
||||
"page": 1,
|
||||
"limit": 20,
|
||||
})
|
||||
}
|
||||
|
||||
// ScrollDocuments scrolls through documents with iterator-style pagination
|
||||
func ScrollDocuments(c *gin.Context) {
|
||||
// TODO: Implement scroll documents logic
|
||||
// Query parameters: cursor, limit, filter, etc.
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"documents": []interface{}{},
|
||||
"cursor": "",
|
||||
"hasMore": false,
|
||||
})
|
||||
}
|
||||
|
||||
// GetDocument gets document details by document ID
|
||||
func GetDocument(c *gin.Context) {
|
||||
// TODO: Implement get document logic
|
||||
// Note: This might need to be implemented based on your document storage structure
|
||||
// as the GraphRag interface doesn't directly provide a GetDocument method
|
||||
docID := c.Param("docID")
|
||||
if docID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "Document ID is required"})
|
||||
return
|
||||
}
|
||||
|
||||
// TODO: Implement actual document retrieval logic
|
||||
// This could involve querying your document storage or getting document metadata
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"docID": docID,
|
||||
"message": "Document details retrieved",
|
||||
// Add actual document fields here when implementing
|
||||
})
|
||||
}
|
||||
|
||||
// RemoveDocs removes documents by IDs
|
||||
func RemoveDocs(c *gin.Context) {
|
||||
// TODO: Implement remove documents logic
|
||||
c.JSON(http.StatusOK, gin.H{"message": "Documents removed"})
|
||||
}
|
||||
|
||||
// Validator interface for request validation
|
||||
type Validator interface {
|
||||
Validate() error
|
||||
|
|
@ -119,378 +166,3 @@ func handleAsync(c *gin.Context, syncHandler func(*gin.Context)) {
|
|||
|
||||
response.RespondWithSuccess(c, response.StatusCreated, gin.H{"job_id": jobid})
|
||||
}
|
||||
|
||||
// AddFile adds a file to a collection
|
||||
func AddFile(c *gin.Context) {
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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
|
||||
_, 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 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": req.DocID,
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusCreated, result)
|
||||
}
|
||||
|
||||
// AddFileAsync adds file to a collection asynchronously
|
||||
func AddFileAsync(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 {
|
||||
return
|
||||
}
|
||||
|
||||
// Validate file and get path
|
||||
_, _, err := validateFileAndGetPath(c, &req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert request to UpsertOptions (just for validation)
|
||||
_, err = getUpsertOptions(c, &req.BaseUpsertRequest)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle async processing
|
||||
handleAsync(c, AddFile)
|
||||
}
|
||||
|
||||
// AddText adds text to a collection
|
||||
func AddText(c *gin.Context) {
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// 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
|
||||
_, 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 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": req.DocID,
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusCreated, result)
|
||||
}
|
||||
|
||||
// AddTextAsync adds text to a collection asynchronously
|
||||
func AddTextAsync(c *gin.Context) {
|
||||
var req AddTextRequest
|
||||
|
||||
// Validate request
|
||||
if err := validateRequest(c, &req); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert request to UpsertOptions (just for validation)
|
||||
_, err := getUpsertOptions(c, &req.BaseUpsertRequest)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle async processing
|
||||
handleAsync(c, AddText)
|
||||
}
|
||||
|
||||
// AddURL adds a URL to a collection
|
||||
func AddURL(c *gin.Context) {
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// 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
|
||||
_, 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 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": req.DocID,
|
||||
}
|
||||
|
||||
response.RespondWithSuccess(c, response.StatusCreated, result)
|
||||
}
|
||||
|
||||
// AddURLAsync adds a URL to a collection asynchronously
|
||||
func AddURLAsync(c *gin.Context) {
|
||||
var req AddURLRequest
|
||||
|
||||
// Validate request
|
||||
if err := validateRequest(c, &req); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Check if kb.Instance is available
|
||||
if !checkKBInstance(c) {
|
||||
return
|
||||
}
|
||||
|
||||
// Convert request to UpsertOptions (just for validation)
|
||||
_, err := getUpsertOptions(c, &req.BaseUpsertRequest)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Handle async processing
|
||||
handleAsync(c, AddURL)
|
||||
}
|
||||
|
||||
// ListDocuments lists documents with pagination
|
||||
func ListDocuments(c *gin.Context) {
|
||||
// TODO: Implement list documents logic
|
||||
// Query parameters for pagination: page, limit, filter, etc.
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"documents": []interface{}{},
|
||||
"total": 0,
|
||||
"page": 1,
|
||||
"limit": 20,
|
||||
})
|
||||
}
|
||||
|
||||
// ScrollDocuments scrolls through documents with iterator-style pagination
|
||||
func ScrollDocuments(c *gin.Context) {
|
||||
// TODO: Implement scroll documents logic
|
||||
// Query parameters: cursor, limit, filter, etc.
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"documents": []interface{}{},
|
||||
"cursor": "",
|
||||
"hasMore": false,
|
||||
})
|
||||
}
|
||||
|
||||
// GetDocument gets document details by document ID
|
||||
func GetDocument(c *gin.Context) {
|
||||
// TODO: Implement get document logic
|
||||
// Note: This might need to be implemented based on your document storage structure
|
||||
// as the GraphRag interface doesn't directly provide a GetDocument method
|
||||
docID := c.Param("docID")
|
||||
if docID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "Document ID is required"})
|
||||
return
|
||||
}
|
||||
|
||||
// TODO: Implement actual document retrieval logic
|
||||
// This could involve querying your document storage or getting document metadata
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"docID": docID,
|
||||
"message": "Document details retrieved",
|
||||
// Add actual document fields here when implementing
|
||||
})
|
||||
}
|
||||
|
||||
// RemoveDocs removes documents by IDs
|
||||
func RemoveDocs(c *gin.Context) {
|
||||
// TODO: Implement remove documents logic
|
||||
c.JSON(http.StatusOK, gin.H{"message": "Documents removed"})
|
||||
}
|
||||
|
|
|
|||
437
openapi/kb/types.go
Normal file
437
openapi/kb/types.go
Normal file
|
|
@ -0,0 +1,437 @@
|
|||
package kb
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/yaoapp/gou/graphrag/types"
|
||||
"github.com/yaoapp/yao/kb"
|
||||
"github.com/yaoapp/yao/kb/providers/factory"
|
||||
kbtypes "github.com/yaoapp/yao/kb/types"
|
||||
)
|
||||
|
||||
/*
|
||||
Usage Examples:
|
||||
|
||||
1. AddFile API (converter will be auto-detected based on file info):
|
||||
{
|
||||
"collection_id": "my_collection",
|
||||
"locale": "en",
|
||||
"file_id": "uploaded_file_123",
|
||||
"chunking": {
|
||||
"provider_id": "__yao.structured",
|
||||
"option_id": "standard"
|
||||
},
|
||||
"embedding": {
|
||||
"provider_id": "__yao.openai",
|
||||
"option_id": "text-embedding-3-small"
|
||||
},
|
||||
"doc_id": "document_001",
|
||||
"metadata": {
|
||||
"source": "research_paper"
|
||||
}
|
||||
}
|
||||
|
||||
2. AddText API with Chinese locale:
|
||||
{
|
||||
"collection_id": "my_collection",
|
||||
"locale": "zh-cn",
|
||||
"text": "这是要处理的文本内容。",
|
||||
"chunking": {
|
||||
"provider_id": "__yao.structured"
|
||||
},
|
||||
"embedding": {
|
||||
"provider_id": "__yao.fastembed",
|
||||
"option_id": "fastembed-chinese"
|
||||
}
|
||||
}
|
||||
|
||||
3. AddSegments API:
|
||||
{
|
||||
"collection_id": "my_collection",
|
||||
"locale": "en",
|
||||
"doc_id": "document_001",
|
||||
"segment_texts": [
|
||||
{"text": "First segment", "metadata": {"page": 1}},
|
||||
{"text": "Second segment", "metadata": {"page": 2}}
|
||||
],
|
||||
"embedding": {
|
||||
"provider_id": "__yao.openai",
|
||||
"option_id": "text-embedding-3-small"
|
||||
}
|
||||
}
|
||||
|
||||
Note:
|
||||
- If no locale is specified, defaults to "en"
|
||||
- If no option_id is specified, the default option from provider configuration will be selected
|
||||
- Providers are loaded based on locale with fallback to "en" if the specified locale is not available
|
||||
- For AddFile API, converter will be auto-detected based on filename and content_type obtained from GetFileInfo(file_id)
|
||||
- ToUpsertOptions() can be called without parameters, or with filename and contentType for converter auto-detection
|
||||
*/
|
||||
|
||||
// ProviderConfig represents a provider configuration that can be specified in two ways:
|
||||
// 1. ProviderID + OptionID (option will be looked up from provider)
|
||||
// 2. ProviderID + Option (option is provided directly)
|
||||
type ProviderConfig struct {
|
||||
ProviderID string `json:"provider_id" binding:"required"`
|
||||
OptionID string `json:"option_id,omitempty"`
|
||||
Option *kbtypes.ProviderOption `json:"option,omitempty"`
|
||||
}
|
||||
|
||||
// BaseUpsertRequest contains common fields for all upsert operations
|
||||
type BaseUpsertRequest struct {
|
||||
// Collection ID - this will be mapped to UpsertOptions.CollectionID
|
||||
CollectionID string `json:"collection_id" binding:"required"`
|
||||
|
||||
// Language/locale for provider selection (defaults to "en")
|
||||
Locale string `json:"locale,omitempty"`
|
||||
|
||||
// Provider configurations
|
||||
Chunking *ProviderConfig `json:"chunking" binding:"required"`
|
||||
Embedding *ProviderConfig `json:"embedding" binding:"required"`
|
||||
Extraction *ProviderConfig `json:"extraction,omitempty"`
|
||||
Fetcher *ProviderConfig `json:"fetcher,omitempty"`
|
||||
Converter *ProviderConfig `json:"converter,omitempty"`
|
||||
|
||||
// Upsert options
|
||||
DocID string `json:"doc_id,omitempty"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// AddFileRequest represents the request for AddFile API
|
||||
type AddFileRequest struct {
|
||||
BaseUpsertRequest
|
||||
FileID string `json:"file_id" binding:"required"`
|
||||
Uploader string `json:"uploader,omitempty"` // The name of the uploader, e.g. "s3", "local", "webdav", etc.
|
||||
}
|
||||
|
||||
// AddTextRequest represents the request for AddText API
|
||||
type AddTextRequest struct {
|
||||
BaseUpsertRequest
|
||||
Text string `json:"text" binding:"required"`
|
||||
}
|
||||
|
||||
// AddURLRequest represents the request for AddURL API
|
||||
type AddURLRequest struct {
|
||||
BaseUpsertRequest
|
||||
URL string `json:"url" binding:"required"`
|
||||
}
|
||||
|
||||
// AddSegmentsRequest represents the request for AddSegments API
|
||||
type AddSegmentsRequest struct {
|
||||
BaseUpsertRequest
|
||||
SegmentTexts []types.SegmentText `json:"segment_texts" binding:"required"`
|
||||
}
|
||||
|
||||
// UpdateSegmentsRequest represents the request for UpdateSegments API
|
||||
type UpdateSegmentsRequest struct {
|
||||
BaseUpsertRequest
|
||||
SegmentTexts []types.SegmentText `json:"segment_texts" binding:"required"`
|
||||
}
|
||||
|
||||
// ProviderOption resolves a ProviderConfig to a *kbtypes.ProviderOption
|
||||
// If OptionID is provided, it looks up the option from the provider
|
||||
// If Option is provided directly, it uses the Option field
|
||||
// If neither is provided, it selects the default option from provider's Options
|
||||
func (config *ProviderConfig) ProviderOption(providerType, locale string) (*kbtypes.ProviderOption, error) {
|
||||
if config == nil {
|
||||
return nil, fmt.Errorf("provider config is required")
|
||||
}
|
||||
|
||||
if config.ProviderID == "" {
|
||||
return nil, fmt.Errorf("provider_id is required")
|
||||
}
|
||||
|
||||
if providerType == "" {
|
||||
return nil, fmt.Errorf("provider_type is required")
|
||||
}
|
||||
|
||||
// If Option is provided directly, use it
|
||||
if config.Option != nil {
|
||||
return config.Option, nil
|
||||
}
|
||||
|
||||
// Get the provider from KB instance
|
||||
if kb.Instance == nil {
|
||||
return nil, fmt.Errorf("KB instance is not initialized")
|
||||
}
|
||||
|
||||
// Default locale to "en" if not provided
|
||||
if locale == "" {
|
||||
locale = "en"
|
||||
}
|
||||
|
||||
// Find the provider using the specified provider type
|
||||
var provider *kbtypes.Provider
|
||||
kbInstance := kb.Instance.(*kb.KnowledgeBase)
|
||||
|
||||
// Get providers of the specific type
|
||||
providers := kbInstance.Providers.GetProviders(providerType, locale)
|
||||
for _, p := range providers {
|
||||
if p.ID == config.ProviderID {
|
||||
provider = p
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if provider == nil {
|
||||
return nil, fmt.Errorf("provider %s not found for locale %s", config.ProviderID, locale)
|
||||
}
|
||||
|
||||
// If OptionID is provided, look it up from the provider
|
||||
if config.OptionID != "" {
|
||||
option, exists := provider.GetOption(config.OptionID)
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("option %s not found in provider %s", config.OptionID, config.ProviderID)
|
||||
}
|
||||
return option, nil
|
||||
}
|
||||
|
||||
// If no option specified, try to find the default option
|
||||
if provider.Options != nil {
|
||||
for _, option := range provider.Options {
|
||||
if option.Default {
|
||||
return option, nil
|
||||
}
|
||||
}
|
||||
// If no default option found but options exist, return the first one
|
||||
if len(provider.Options) > 0 {
|
||||
return provider.Options[0], nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("no option specified and no default option found for provider %s", config.ProviderID)
|
||||
}
|
||||
|
||||
// ToUpsertOptions converts BaseUpsertRequest to types.UpsertOptions
|
||||
// Optional parameters: filename, contentType (for converter auto-detection)
|
||||
func (r *BaseUpsertRequest) ToUpsertOptions(fileInfo ...string) (*types.UpsertOptions, error) {
|
||||
var filename, contentType string
|
||||
if len(fileInfo) >= 1 {
|
||||
filename = fileInfo[0]
|
||||
}
|
||||
if len(fileInfo) >= 2 {
|
||||
contentType = fileInfo[1]
|
||||
}
|
||||
|
||||
// Default locale to "en" if not specified
|
||||
locale := r.Locale
|
||||
if locale == "" {
|
||||
locale = "en"
|
||||
}
|
||||
|
||||
options := &types.UpsertOptions{
|
||||
CollectionID: r.CollectionID, // Collection ID maps to CollectionID
|
||||
DocID: r.DocID,
|
||||
Metadata: r.Metadata,
|
||||
}
|
||||
|
||||
// Resolve and create chunking provider
|
||||
chunkingOption, err := r.Chunking.ProviderOption("chunking", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve chunking provider: %w", err)
|
||||
}
|
||||
|
||||
chunking, err := factory.MakeChunking(r.Chunking.ProviderID, chunkingOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create chunking provider: %w", err)
|
||||
}
|
||||
options.Chunking = chunking
|
||||
|
||||
// Get chunking options
|
||||
chunkingOpts, err := factory.ChunkingOptions(r.Chunking.ProviderID, chunkingOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get chunking options: %w", err)
|
||||
}
|
||||
options.ChunkingOptions = chunkingOpts
|
||||
|
||||
// Resolve and create embedding provider
|
||||
embeddingOption, err := r.Embedding.ProviderOption("embedding", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve embedding provider: %w", err)
|
||||
}
|
||||
|
||||
embedding, err := factory.MakeEmbedding(r.Embedding.ProviderID, embeddingOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create embedding provider: %w", err)
|
||||
}
|
||||
options.Embedding = embedding
|
||||
|
||||
// Optional providers
|
||||
if r.Extraction != nil {
|
||||
extractionOption, err := r.Extraction.ProviderOption("extraction", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve extraction provider: %w", err)
|
||||
}
|
||||
|
||||
extraction, err := factory.MakeExtraction(r.Extraction.ProviderID, extractionOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create extraction provider: %w", err)
|
||||
}
|
||||
options.Extraction = extraction
|
||||
}
|
||||
|
||||
if r.Fetcher != nil {
|
||||
fetcherOption, err := r.Fetcher.ProviderOption("fetcher", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve fetcher provider: %w", err)
|
||||
}
|
||||
|
||||
fetcher, err := factory.MakeFetcher(r.Fetcher.ProviderID, fetcherOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create fetcher provider: %w", err)
|
||||
}
|
||||
options.Fetcher = fetcher
|
||||
}
|
||||
|
||||
// Handle converter - auto-detect if not specified
|
||||
if r.Converter != nil {
|
||||
// User specified converter
|
||||
converterOption, err := r.Converter.ProviderOption("converter", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve converter provider: %w", err)
|
||||
}
|
||||
|
||||
converter, err := factory.MakeConverter(r.Converter.ProviderID, converterOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create converter provider: %w", err)
|
||||
}
|
||||
options.Converter = converter
|
||||
} else if filename != "" || contentType != "" {
|
||||
// Auto-detect converter based on filename and content type
|
||||
matched, converterID, err := factory.AutoDetectConverter(filename, contentType)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to auto-detect converter: %w", err)
|
||||
}
|
||||
|
||||
if matched {
|
||||
// Find the provider to get default option
|
||||
converterConfig := &ProviderConfig{
|
||||
ProviderID: converterID,
|
||||
}
|
||||
|
||||
converterOption, err := converterConfig.ProviderOption("converter", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve auto-detected converter provider: %w", err)
|
||||
}
|
||||
|
||||
converter, err := factory.MakeConverter(converterID, converterOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create auto-detected converter provider: %w", err)
|
||||
}
|
||||
options.Converter = converter
|
||||
}
|
||||
}
|
||||
|
||||
return options, nil
|
||||
}
|
||||
|
||||
// Validate validates the common fields
|
||||
func (r *BaseUpsertRequest) Validate() error {
|
||||
if r.CollectionID == "" {
|
||||
return fmt.Errorf("collection_id is required")
|
||||
}
|
||||
if r.Chunking == nil {
|
||||
return fmt.Errorf("chunking provider is required")
|
||||
}
|
||||
if r.Embedding == nil {
|
||||
return fmt.Errorf("embedding provider is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddFileRequest fields
|
||||
func (r *AddFileRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if r.FileID == "" {
|
||||
return fmt.Errorf("file_id is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddTextRequest fields
|
||||
func (r *AddTextRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if r.Text == "" {
|
||||
return fmt.Errorf("text is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddURLRequest fields
|
||||
func (r *AddURLRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if r.URL == "" {
|
||||
return fmt.Errorf("url is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddSegmentsRequest fields
|
||||
func (r *AddSegmentsRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(r.SegmentTexts) == 0 {
|
||||
return fmt.Errorf("segment_texts is required")
|
||||
}
|
||||
if r.DocID == "" {
|
||||
return fmt.Errorf("doc_id is required for AddSegments operation")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the UpdateSegmentsRequest fields
|
||||
func (r *UpdateSegmentsRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(r.SegmentTexts) == 0 {
|
||||
return fmt.Errorf("segment_texts is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddBaseFields adds common fields from BaseUpsertRequest to data map
|
||||
func (r *BaseUpsertRequest) AddBaseFields(data map[string]interface{}) {
|
||||
if r.Locale != "" {
|
||||
data["locale"] = r.Locale
|
||||
}
|
||||
if r.DocID != "" {
|
||||
data["document_id"] = r.DocID
|
||||
}
|
||||
if r.Metadata != nil {
|
||||
data["tags"] = r.Metadata
|
||||
}
|
||||
|
||||
// Add provider configurations
|
||||
if r.Converter != nil {
|
||||
data["converter_provider_id"] = r.Converter.ProviderID
|
||||
if r.Converter.Option != nil {
|
||||
data["converter_properties"] = r.Converter.Option.Properties
|
||||
}
|
||||
}
|
||||
if r.Fetcher != nil {
|
||||
data["fetcher_provider_id"] = r.Fetcher.ProviderID
|
||||
if r.Fetcher.Option != nil {
|
||||
data["fetcher_properties"] = r.Fetcher.Option.Properties
|
||||
}
|
||||
}
|
||||
if r.Chunking != nil {
|
||||
data["chunking_provider_id"] = r.Chunking.ProviderID
|
||||
if r.Chunking.Option != nil {
|
||||
data["chunking_properties"] = r.Chunking.Option.Properties
|
||||
}
|
||||
}
|
||||
if r.Extraction != nil {
|
||||
data["extraction_provider_id"] = r.Extraction.ProviderID
|
||||
if r.Extraction.Option != nil {
|
||||
data["extraction_properties"] = r.Extraction.Option.Properties
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -4,402 +4,10 @@ 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"
|
||||
)
|
||||
|
||||
/*
|
||||
Usage Examples:
|
||||
|
||||
1. AddFile API (converter will be auto-detected based on file info):
|
||||
{
|
||||
"collection_id": "my_collection",
|
||||
"locale": "en",
|
||||
"file_id": "uploaded_file_123",
|
||||
"chunking": {
|
||||
"provider_id": "__yao.structured",
|
||||
"option_id": "standard"
|
||||
},
|
||||
"embedding": {
|
||||
"provider_id": "__yao.openai",
|
||||
"option_id": "text-embedding-3-small"
|
||||
},
|
||||
"doc_id": "document_001",
|
||||
"metadata": {
|
||||
"source": "research_paper"
|
||||
}
|
||||
}
|
||||
|
||||
2. AddText API with Chinese locale:
|
||||
{
|
||||
"collection_id": "my_collection",
|
||||
"locale": "zh-cn",
|
||||
"text": "这是要处理的文本内容。",
|
||||
"chunking": {
|
||||
"provider_id": "__yao.structured"
|
||||
},
|
||||
"embedding": {
|
||||
"provider_id": "__yao.fastembed",
|
||||
"option_id": "fastembed-chinese"
|
||||
}
|
||||
}
|
||||
|
||||
3. AddSegments API:
|
||||
{
|
||||
"collection_id": "my_collection",
|
||||
"locale": "en",
|
||||
"doc_id": "document_001",
|
||||
"segment_texts": [
|
||||
{"text": "First segment", "metadata": {"page": 1}},
|
||||
{"text": "Second segment", "metadata": {"page": 2}}
|
||||
],
|
||||
"embedding": {
|
||||
"provider_id": "__yao.openai",
|
||||
"option_id": "text-embedding-3-small"
|
||||
}
|
||||
}
|
||||
|
||||
Note:
|
||||
- If no locale is specified, defaults to "en"
|
||||
- If no option_id is specified, the default option from provider configuration will be selected
|
||||
- Providers are loaded based on locale with fallback to "en" if the specified locale is not available
|
||||
- For AddFile API, converter will be auto-detected based on filename and content_type obtained from GetFileInfo(file_id)
|
||||
- ToUpsertOptions() can be called without parameters, or with filename and contentType for converter auto-detection
|
||||
*/
|
||||
|
||||
// ProviderConfig represents a provider configuration that can be specified in two ways:
|
||||
// 1. ProviderID + OptionID (option will be looked up from provider)
|
||||
// 2. ProviderID + Option (option is provided directly)
|
||||
type ProviderConfig struct {
|
||||
ProviderID string `json:"provider_id" binding:"required"`
|
||||
OptionID string `json:"option_id,omitempty"`
|
||||
Option *kbtypes.ProviderOption `json:"option,omitempty"`
|
||||
}
|
||||
|
||||
// BaseUpsertRequest contains common fields for all upsert operations
|
||||
type BaseUpsertRequest struct {
|
||||
// Collection ID - this will be mapped to UpsertOptions.CollectionID
|
||||
CollectionID string `json:"collection_id" binding:"required"`
|
||||
|
||||
// Language/locale for provider selection (defaults to "en")
|
||||
Locale string `json:"locale,omitempty"`
|
||||
|
||||
// Provider configurations
|
||||
Chunking *ProviderConfig `json:"chunking" binding:"required"`
|
||||
Embedding *ProviderConfig `json:"embedding" binding:"required"`
|
||||
Extraction *ProviderConfig `json:"extraction,omitempty"`
|
||||
Fetcher *ProviderConfig `json:"fetcher,omitempty"`
|
||||
Converter *ProviderConfig `json:"converter,omitempty"`
|
||||
|
||||
// Upsert options
|
||||
DocID string `json:"doc_id,omitempty"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// AddFileRequest represents the request for AddFile API
|
||||
type AddFileRequest struct {
|
||||
BaseUpsertRequest
|
||||
FileID string `json:"file_id" binding:"required"`
|
||||
Uploader string `json:"uploader,omitempty"` // The name of the uploader, e.g. "s3", "local", "webdav", etc.
|
||||
}
|
||||
|
||||
// AddTextRequest represents the request for AddText API
|
||||
type AddTextRequest struct {
|
||||
BaseUpsertRequest
|
||||
Text string `json:"text" binding:"required"`
|
||||
}
|
||||
|
||||
// AddURLRequest represents the request for AddURL API
|
||||
type AddURLRequest struct {
|
||||
BaseUpsertRequest
|
||||
URL string `json:"url" binding:"required"`
|
||||
}
|
||||
|
||||
// AddSegmentsRequest represents the request for AddSegments API
|
||||
type AddSegmentsRequest struct {
|
||||
BaseUpsertRequest
|
||||
SegmentTexts []types.SegmentText `json:"segment_texts" binding:"required"`
|
||||
}
|
||||
|
||||
// UpdateSegmentsRequest represents the request for UpdateSegments API
|
||||
type UpdateSegmentsRequest struct {
|
||||
BaseUpsertRequest
|
||||
SegmentTexts []types.SegmentText `json:"segment_texts" binding:"required"`
|
||||
}
|
||||
|
||||
// resolveProviderOption resolves a ProviderConfig to a *kbtypes.ProviderOption
|
||||
// If OptionID is provided, it looks up the option from the provider
|
||||
// If Option is provided directly, it uses the Option field
|
||||
// If neither is provided, it selects the default option from provider's Options
|
||||
func resolveProviderOption(config *ProviderConfig, providerType, locale string) (*kbtypes.ProviderOption, error) {
|
||||
if config == nil {
|
||||
return nil, fmt.Errorf("provider config is required")
|
||||
}
|
||||
|
||||
if config.ProviderID == "" {
|
||||
return nil, fmt.Errorf("provider_id is required")
|
||||
}
|
||||
|
||||
if providerType == "" {
|
||||
return nil, fmt.Errorf("provider_type is required")
|
||||
}
|
||||
|
||||
// If Option is provided directly, use it
|
||||
if config.Option != nil {
|
||||
return config.Option, nil
|
||||
}
|
||||
|
||||
// Get the provider from KB instance
|
||||
if kb.Instance == nil {
|
||||
return nil, fmt.Errorf("KB instance is not initialized")
|
||||
}
|
||||
|
||||
// Default locale to "en" if not provided
|
||||
if locale == "" {
|
||||
locale = "en"
|
||||
}
|
||||
|
||||
// Find the provider using the specified provider type
|
||||
var provider *kbtypes.Provider
|
||||
kbInstance := kb.Instance.(*kb.KnowledgeBase)
|
||||
|
||||
// Get providers of the specific type
|
||||
providers := kbInstance.Providers.GetProviders(providerType, locale)
|
||||
for _, p := range providers {
|
||||
if p.ID == config.ProviderID {
|
||||
provider = p
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if provider == nil {
|
||||
return nil, fmt.Errorf("provider %s not found for locale %s", config.ProviderID, locale)
|
||||
}
|
||||
|
||||
// If OptionID is provided, look it up from the provider
|
||||
if config.OptionID != "" {
|
||||
option, exists := provider.GetOption(config.OptionID)
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("option %s not found in provider %s", config.OptionID, config.ProviderID)
|
||||
}
|
||||
return option, nil
|
||||
}
|
||||
|
||||
// If no option specified, try to find the default option
|
||||
if provider.Options != nil {
|
||||
for _, option := range provider.Options {
|
||||
if option.Default {
|
||||
return option, nil
|
||||
}
|
||||
}
|
||||
// If no default option found but options exist, return the first one
|
||||
if len(provider.Options) > 0 {
|
||||
return provider.Options[0], nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("no option specified and no default option found for provider %s", config.ProviderID)
|
||||
}
|
||||
|
||||
// ToUpsertOptions converts BaseUpsertRequest to types.UpsertOptions
|
||||
// Optional parameters: filename, contentType (for converter auto-detection)
|
||||
func (r *BaseUpsertRequest) ToUpsertOptions(fileInfo ...string) (*types.UpsertOptions, error) {
|
||||
var filename, contentType string
|
||||
if len(fileInfo) >= 1 {
|
||||
filename = fileInfo[0]
|
||||
}
|
||||
if len(fileInfo) >= 2 {
|
||||
contentType = fileInfo[1]
|
||||
}
|
||||
|
||||
// Default locale to "en" if not specified
|
||||
locale := r.Locale
|
||||
if locale == "" {
|
||||
locale = "en"
|
||||
}
|
||||
|
||||
options := &types.UpsertOptions{
|
||||
CollectionID: r.CollectionID, // Collection ID maps to CollectionID
|
||||
DocID: r.DocID,
|
||||
Metadata: r.Metadata,
|
||||
}
|
||||
|
||||
// Resolve and create chunking provider
|
||||
chunkingOption, err := resolveProviderOption(r.Chunking, "chunking", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve chunking provider: %w", err)
|
||||
}
|
||||
|
||||
chunking, err := factory.MakeChunking(r.Chunking.ProviderID, chunkingOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create chunking provider: %w", err)
|
||||
}
|
||||
options.Chunking = chunking
|
||||
|
||||
// Get chunking options
|
||||
chunkingOpts, err := factory.ChunkingOptions(r.Chunking.ProviderID, chunkingOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get chunking options: %w", err)
|
||||
}
|
||||
options.ChunkingOptions = chunkingOpts
|
||||
|
||||
// Resolve and create embedding provider
|
||||
embeddingOption, err := resolveProviderOption(r.Embedding, "embedding", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve embedding provider: %w", err)
|
||||
}
|
||||
|
||||
embedding, err := factory.MakeEmbedding(r.Embedding.ProviderID, embeddingOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create embedding provider: %w", err)
|
||||
}
|
||||
options.Embedding = embedding
|
||||
|
||||
// Optional providers
|
||||
if r.Extraction != nil {
|
||||
extractionOption, err := resolveProviderOption(r.Extraction, "extraction", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve extraction provider: %w", err)
|
||||
}
|
||||
|
||||
extraction, err := factory.MakeExtraction(r.Extraction.ProviderID, extractionOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create extraction provider: %w", err)
|
||||
}
|
||||
options.Extraction = extraction
|
||||
}
|
||||
|
||||
if r.Fetcher != nil {
|
||||
fetcherOption, err := resolveProviderOption(r.Fetcher, "fetcher", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve fetcher provider: %w", err)
|
||||
}
|
||||
|
||||
fetcher, err := factory.MakeFetcher(r.Fetcher.ProviderID, fetcherOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create fetcher provider: %w", err)
|
||||
}
|
||||
options.Fetcher = fetcher
|
||||
}
|
||||
|
||||
// Handle converter - auto-detect if not specified
|
||||
if r.Converter != nil {
|
||||
// User specified converter
|
||||
converterOption, err := resolveProviderOption(r.Converter, "converter", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve converter provider: %w", err)
|
||||
}
|
||||
|
||||
converter, err := factory.MakeConverter(r.Converter.ProviderID, converterOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create converter provider: %w", err)
|
||||
}
|
||||
options.Converter = converter
|
||||
} else if filename != "" || contentType != "" {
|
||||
// Auto-detect converter based on filename and content type
|
||||
matched, converterID, err := factory.AutoDetectConverter(filename, contentType)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to auto-detect converter: %w", err)
|
||||
}
|
||||
|
||||
if matched {
|
||||
// Find the provider to get default option
|
||||
converterConfig := &ProviderConfig{
|
||||
ProviderID: converterID,
|
||||
}
|
||||
|
||||
converterOption, err := resolveProviderOption(converterConfig, "converter", locale)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to resolve auto-detected converter provider: %w", err)
|
||||
}
|
||||
|
||||
converter, err := factory.MakeConverter(converterID, converterOption)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create auto-detected converter provider: %w", err)
|
||||
}
|
||||
options.Converter = converter
|
||||
}
|
||||
}
|
||||
|
||||
return options, nil
|
||||
}
|
||||
|
||||
// Validate validates the common fields
|
||||
func (r *BaseUpsertRequest) Validate() error {
|
||||
if r.CollectionID == "" {
|
||||
return fmt.Errorf("collection_id is required")
|
||||
}
|
||||
if r.Chunking == nil {
|
||||
return fmt.Errorf("chunking provider is required")
|
||||
}
|
||||
if r.Embedding == nil {
|
||||
return fmt.Errorf("embedding provider is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddFileRequest fields
|
||||
func (r *AddFileRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if r.FileID == "" {
|
||||
return fmt.Errorf("file_id is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddTextRequest fields
|
||||
func (r *AddTextRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if r.Text == "" {
|
||||
return fmt.Errorf("text is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddURLRequest fields
|
||||
func (r *AddURLRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if r.URL == "" {
|
||||
return fmt.Errorf("url is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the AddSegmentsRequest fields
|
||||
func (r *AddSegmentsRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(r.SegmentTexts) == 0 {
|
||||
return fmt.Errorf("segment_texts is required")
|
||||
}
|
||||
if r.DocID == "" {
|
||||
return fmt.Errorf("doc_id is required for AddSegments operation")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Validate validates the UpdateSegmentsRequest fields
|
||||
func (r *UpdateSegmentsRequest) Validate() error {
|
||||
if err := r.BaseUpsertRequest.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(r.SegmentTexts) == 0 {
|
||||
return fmt.Errorf("segment_texts is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PrepareCreateCollection prepares CreateCollection request and database data
|
||||
func PrepareCreateCollection(c *gin.Context) (*CreateCollectionRequest, map[string]interface{}, error) {
|
||||
var req CreateCollectionRequest
|
||||
|
|
@ -509,7 +117,7 @@ func PrepareAddFile(c *gin.Context) (*AddFileRequest, map[string]interface{}, er
|
|||
"size": int64(fileInfo.Bytes),
|
||||
}
|
||||
|
||||
addBaseRequestFields(data, &req.BaseUpsertRequest)
|
||||
req.BaseUpsertRequest.AddBaseFields(data)
|
||||
addContextFields(c, data)
|
||||
|
||||
return &req, data, nil
|
||||
|
|
@ -547,7 +155,7 @@ func PrepareAddText(c *gin.Context) (*AddTextRequest, map[string]interface{}, er
|
|||
}
|
||||
}
|
||||
|
||||
addBaseRequestFields(data, &req.BaseUpsertRequest)
|
||||
req.BaseUpsertRequest.AddBaseFields(data)
|
||||
addContextFields(c, data)
|
||||
|
||||
return &req, data, nil
|
||||
|
|
@ -585,51 +193,12 @@ func PrepareAddURL(c *gin.Context) (*AddURLRequest, map[string]interface{}, erro
|
|||
}
|
||||
}
|
||||
|
||||
addBaseRequestFields(data, &req.BaseUpsertRequest)
|
||||
req.BaseUpsertRequest.AddBaseFields(data)
|
||||
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["extraction_provider_id"] = req.Extraction.ProviderID
|
||||
if req.Extraction.Option != nil {
|
||||
data["extraction_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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue