mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-05-28 11:43:06 +08:00
af85aa9c7b
### What problem does this PR solve? Implement Elasticsearch functions in GO (except for Search) ### Type of change - [x] Refactoring
545 lines
14 KiB
Go
545 lines
14 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package handler
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"os"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/engine"
|
|
"ragflow/internal/service"
|
|
)
|
|
|
|
// TenantHandler tenant handler
|
|
type TenantHandler struct {
|
|
tenantService *service.TenantService
|
|
userService *service.UserService
|
|
kbService *service.KnowledgebaseService
|
|
}
|
|
|
|
// NewTenantHandler create tenant handler
|
|
func NewTenantHandler(tenantService *service.TenantService, userService *service.UserService, kbService *service.KnowledgebaseService) *TenantHandler {
|
|
return &TenantHandler{
|
|
tenantService: tenantService,
|
|
userService: userService,
|
|
kbService: kbService,
|
|
}
|
|
}
|
|
|
|
func (h *TenantHandler) GetModels(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
defaultModels, err := h.tenantService.ListTenantDefaultModels(user.ID)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeExceptionError,
|
|
"message": err.Error(),
|
|
"data": false,
|
|
})
|
|
return
|
|
}
|
|
|
|
if defaultModels == nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeDataError,
|
|
"message": "No default models",
|
|
"data": nil,
|
|
})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": defaultModels,
|
|
})
|
|
}
|
|
|
|
type SetModelRequest struct {
|
|
ModelProvider string `json:"model_provider"`
|
|
ModelInstance string `json:"model_instance"`
|
|
ModelName string `json:"model_name"`
|
|
ModelType string `json:"model_type" binding:"required"`
|
|
}
|
|
|
|
func (h *TenantHandler) SetModels(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
// Parse request body (same as Python get_request_json())
|
|
var req SetModelRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": common.CodeBadRequest,
|
|
"data": nil,
|
|
"message": "Invalid request body: " + err.Error(),
|
|
})
|
|
return
|
|
}
|
|
|
|
err := h.tenantService.SetTenantDefaultModels(user.ID, req.ModelProvider, req.ModelInstance, req.ModelName, req.ModelType)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeExceptionError,
|
|
"message": err.Error(),
|
|
"data": false,
|
|
})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": nil,
|
|
})
|
|
}
|
|
|
|
// TenantInfo get tenant information
|
|
// @Summary Get Tenant Information
|
|
// @Description Get current user's tenant information (owner tenant)
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/user/tenant_info [get]
|
|
func (h *TenantHandler) TenantInfo(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
tenantInfo, err := h.tenantService.GetTenantInfo(user.ID)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeExceptionError,
|
|
"message": err.Error(),
|
|
"data": false,
|
|
})
|
|
return
|
|
}
|
|
|
|
if tenantInfo == nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeDataError,
|
|
"message": "Tenant not found!",
|
|
"data": false,
|
|
})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": tenantInfo,
|
|
})
|
|
}
|
|
|
|
// TenantList get tenant list for current user
|
|
// @Summary Get Tenant List
|
|
// @Description Get all tenants that the current user belongs to
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/tenant/list [get]
|
|
func (h *TenantHandler) TenantList(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
tenantList, err := h.tenantService.GetTenantList(user.ID)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeExceptionError,
|
|
"message": err.Error(),
|
|
"data": false,
|
|
})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": tenantList,
|
|
})
|
|
}
|
|
|
|
// CreateMetadataStore handles the create metadata store request
|
|
// @Summary Create Metadata Store
|
|
// @Description Create the metadata store for a tenant
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/tenant/metadata_store [post]
|
|
func (h *TenantHandler) CreateMetadataStore(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
// Use user.ID as tenant ID (user IS the tenant in user mode)
|
|
tenantID := user.ID
|
|
|
|
code, err := h.tenantService.CreateMetadataStore(tenantID)
|
|
if err != nil {
|
|
jsonError(c, code, err.Error())
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": nil,
|
|
})
|
|
}
|
|
|
|
// DeleteMetadataStore handles the delete metadata store request
|
|
// @Summary Delete Metadata Store
|
|
// @Description Delete the metadata store for a tenant
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/tenant/metadata_store [delete]
|
|
func (h *TenantHandler) DeleteMetadataStore(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
// Use user.ID as tenant ID (user IS the tenant in user mode)
|
|
tenantID := user.ID
|
|
|
|
code, err := h.tenantService.DeleteMetadataStore(tenantID)
|
|
if err != nil {
|
|
jsonError(c, code, err.Error())
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": nil,
|
|
})
|
|
}
|
|
|
|
// CreateChunkTableRequest represents the request for creating a chunk table
|
|
type CreateChunkTableRequest struct {
|
|
KBID string `json:"kb_id" binding:"required"`
|
|
VectorSize int `json:"vector_size" binding:"required"`
|
|
}
|
|
|
|
// CreateChunkStore handles the create chunk store request
|
|
// @Summary Create Chunk Store
|
|
// @Description Create the chunk store for a knowledge base
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Param request body CreateChunkTableRequest true "create chunk store request"
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/tenant/chunk_store [post]
|
|
func (h *TenantHandler) CreateChunkStore(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
var req CreateChunkTableRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
jsonError(c, common.CodeDataError, err.Error())
|
|
return
|
|
}
|
|
|
|
// Check authorization - user must have access to this kb
|
|
if !h.kbService.Accessible(req.KBID, user.ID) {
|
|
jsonError(c, common.CodeAuthenticationError, "No authorization.")
|
|
return
|
|
}
|
|
|
|
serviceReq := &service.CreateDatasetTableRequest{
|
|
KBID: req.KBID,
|
|
VectorSize: req.VectorSize,
|
|
}
|
|
result, code, err := h.tenantService.CreateChunkStore(serviceReq)
|
|
if err != nil {
|
|
jsonError(c, code, err.Error())
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": result,
|
|
})
|
|
}
|
|
|
|
// DeleteChunkTableRequest represents the request for deleting a chunk table
|
|
type DeleteChunkTableRequest struct {
|
|
KBID string `json:"kb_id" binding:"required"`
|
|
}
|
|
|
|
// DeleteChunkStore handles the delete chunk store request
|
|
// @Summary Delete Chunk Store
|
|
// @Description Delete the chunk store for a knowledge base
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Param request body DeleteChunkTableRequest true "delete chunk store request"
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/tenant/chunk_store [delete]
|
|
func (h *TenantHandler) DeleteChunkStore(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
var req DeleteChunkTableRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
jsonError(c, common.CodeDataError, err.Error())
|
|
return
|
|
}
|
|
|
|
// Check authorization
|
|
if !h.kbService.Accessible(req.KBID, user.ID) {
|
|
jsonError(c, common.CodeAuthenticationError, "No authorization.")
|
|
return
|
|
}
|
|
|
|
code, err := h.tenantService.DeleteChunkStore(req.KBID)
|
|
if err != nil {
|
|
jsonError(c, code, err.Error())
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": common.CodeSuccess,
|
|
"message": "success",
|
|
"data": nil,
|
|
})
|
|
}
|
|
|
|
// InsertChunksFromFileRequest request for inserting chunks from file
|
|
type InsertChunksFromFileRequest struct {
|
|
FilePath string `json:"file_path" binding:"required"`
|
|
}
|
|
|
|
// @Summary Insert chunks into dataset from JSON file
|
|
// @Description Internal: Insert chunks into dataset table from a JSON file
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Param request body InsertChunksFromFileRequest true "insert chunks request"
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/tenant/insert_chunks_from_file [post]
|
|
func (h *TenantHandler) InsertChunksFromFile(c *gin.Context) {
|
|
_, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
var req InsertChunksFromFileRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": err.Error(),
|
|
})
|
|
return
|
|
}
|
|
|
|
if req.FilePath == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "file_path is required",
|
|
})
|
|
return
|
|
}
|
|
|
|
// Read the JSON file
|
|
data, err := os.ReadFile(req.FilePath)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "failed to read file: " + err.Error(),
|
|
})
|
|
return
|
|
}
|
|
|
|
// Parse JSON - format: {"index_name"/"table_name": ..., "knowledgebase_id": ..., "chunks": [...]}
|
|
var debugFormat struct {
|
|
IndexName string `json:"index_name"`
|
|
TableName string `json:"table_name"`
|
|
KnowledgebaseID string `json:"knowledgebase_id"`
|
|
Chunks []map[string]interface{} `json:"chunks"`
|
|
}
|
|
|
|
if err := json.Unmarshal(data, &debugFormat); err != nil || debugFormat.Chunks == nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "invalid JSON format: expected {\"index_name\"/\"table_name\": ..., \"knowledgebase_id\": ..., \"chunks\": [...]}",
|
|
})
|
|
return
|
|
}
|
|
|
|
if len(debugFormat.Chunks) == 0 {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "no chunks found in file",
|
|
})
|
|
return
|
|
}
|
|
|
|
// Support both index_name (ES) and table_name (Infinity) in JSON
|
|
indexName := debugFormat.IndexName
|
|
if indexName == "" {
|
|
indexName = debugFormat.TableName
|
|
}
|
|
|
|
// Get the document engine and insert
|
|
docEngine := engine.Get()
|
|
result, err := docEngine.InsertChunks(c.Request.Context(), debugFormat.Chunks, indexName, debugFormat.KnowledgebaseID)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{
|
|
"code": 500,
|
|
"message": "failed to insert into dataset: " + err.Error(),
|
|
})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": 0,
|
|
"data": result,
|
|
"message": "success",
|
|
})
|
|
}
|
|
|
|
// InsertMetadataFromFileRequest request for inserting metadata from file
|
|
type InsertMetadataFromFileRequest struct {
|
|
FilePath string `json:"file_path" binding:"required"`
|
|
}
|
|
|
|
// @Summary Insert document metadata from JSON file
|
|
// @Description Internal: Insert metadata into tenant's metadata table from a JSON file
|
|
// @Tags tenants
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Security ApiKeyAuth
|
|
// @Param request body InsertMetadataFromFileRequest true "insert metadata request"
|
|
// @Success 200 {object} map[string]interface{}
|
|
// @Router /v1/tenant/insert_metadata_from_file [post]
|
|
func (h *TenantHandler) InsertMetadataFromFile(c *gin.Context) {
|
|
user, errorCode, errorMessage := GetUser(c)
|
|
if errorCode != common.CodeSuccess {
|
|
jsonError(c, errorCode, errorMessage)
|
|
return
|
|
}
|
|
|
|
var req InsertMetadataFromFileRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": err.Error(),
|
|
})
|
|
return
|
|
}
|
|
|
|
if req.FilePath == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "file_path is required",
|
|
})
|
|
return
|
|
}
|
|
|
|
// Read the JSON file
|
|
data, err := os.ReadFile(req.FilePath)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "failed to read file: " + err.Error(),
|
|
})
|
|
return
|
|
}
|
|
|
|
// Parse JSON - format: {"chunks": [...]}
|
|
var inputFormat struct {
|
|
Chunks []map[string]interface{} `json:"chunks"`
|
|
}
|
|
|
|
if err := json.Unmarshal(data, &inputFormat); err != nil || inputFormat.Chunks == nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "invalid JSON format: expected {\"chunks\": [...]}",
|
|
})
|
|
return
|
|
}
|
|
|
|
if len(inputFormat.Chunks) == 0 {
|
|
c.JSON(http.StatusBadRequest, gin.H{
|
|
"code": 400,
|
|
"message": "no chunks found in file",
|
|
})
|
|
return
|
|
}
|
|
|
|
// Use user.ID as tenant ID (user IS the tenant in user mode)
|
|
tenantID := user.ID
|
|
|
|
// Get the document engine and insert
|
|
docEngine := engine.Get()
|
|
result, err := docEngine.InsertMetadata(c.Request.Context(), inputFormat.Chunks, tenantID)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{
|
|
"code": 500,
|
|
"message": "failed to insert metadata: " + err.Error(),
|
|
})
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": 0,
|
|
"data": result,
|
|
"message": "success",
|
|
})
|
|
}
|