package kb import ( "net/http" "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/yao/attachment" "github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/openapi/response" ) // Document Management Handlers // Validator interface for request validation type Validator interface { Validate() error } // validateRequest validates a request by parsing JSON and calling Validate() func validateRequest[T Validator](c *gin.Context, req T) error { // Parse and bind JSON request if err := c.ShouldBindJSON(req); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Invalid request format: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return err } // Validate request if err := req.Validate(); err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return err } return nil } // checkKBInstance checks if kb.Instance is available func checkKBInstance(c *gin.Context) bool { if kb.Instance == nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Knowledge base not initialized", } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return false } return true } // getUpsertOptions converts BaseUpsertRequest to UpsertOptions with optional file info func getUpsertOptions(c *gin.Context, req *BaseUpsertRequest, fileInfo ...string) (*types.UpsertOptions, error) { upsertOptions, err := req.ToUpsertOptions(fileInfo...) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Failed to convert request to upsert options: " + err.Error(), } response.RespondWithError(c, response.StatusBadRequest, errorResp) return nil, err } return upsertOptions, nil } // validateFileAndGetPath validates file manager, file existence and gets local path func validateFileAndGetPath(c *gin.Context, req *AddFileRequest) (string, string, error) { // Get file manager m, ok := attachment.Managers[req.Uploader] if !ok { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Invalid uploader: " + req.Uploader + " not found", } response.RespondWithError(c, response.StatusNotFound, errorResp) return "", "", response.ErrInvalidRequest } // Check if the file exists exists := m.Exists(c.Request.Context(), req.FileID) if !exists { errorResp := &response.ErrorResponse{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "File not found: " + req.FileID, } response.RespondWithError(c, response.StatusNotFound, errorResp) return "", "", response.ErrInvalidRequest } // Get the options of the manager path, contentType, err := m.LocalPath(c.Request.Context(), req.FileID) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to get local path: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return "", "", err } return path, contentType, nil } // handleAsync handles async processing for any handler function func handleAsync(c *gin.Context, syncHandler func(*gin.Context)) { jobid := uuid.New().String() // temporary solution to handle async operations ( TODO: use job queue ) go func() { syncHandler(c) }() response.RespondWithSuccess(c, response.StatusCreated, gin.H{"job_id": jobid}) } // 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 { return } // Validate file and get path path, contentType, err := validateFileAndGetPath(c, &req) if err != nil { return } // Convert request to UpsertOptions upsertOptions, err := getUpsertOptions(c, &req.BaseUpsertRequest, path, contentType) if err != nil { 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) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to upsert file: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Return success response result := gin.H{ "message": "File added successfully", "collection_id": req.CollectionID, "file_id": req.FileID, "doc_id": 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) { 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 upsertOptions, err := getUpsertOptions(c, &req.BaseUpsertRequest) if err != nil { return } // Perform upsert operation with text docID, err := kb.Instance.AddText(c.Request.Context(), req.Text, upsertOptions) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to upsert text: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Return success response result := gin.H{ "message": "Text added successfully", "collection_id": req.CollectionID, "doc_id": 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) { 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 upsertOptions, err := getUpsertOptions(c, &req.BaseUpsertRequest) if err != nil { return } // Perform upsert operation with URL docID, err := kb.Instance.AddURL(c.Request.Context(), req.URL, upsertOptions) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Failed to upsert URL: " + err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return } // Return success response result := gin.H{ "message": "URL added successfully", "collection_id": req.CollectionID, "url": req.URL, "doc_id": 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"}) }