zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-29 20:19:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: add context, part6 (#17399)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-27 11:09:17 +08:00
committed by GitHub
parent 53afc32349
commit 9d4847beaf
14 changed files with 143 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
internal/admin/handler.go
View File

@@ -258,8 +258,9 @@ func (h *Handler) CreateUser(c *gin.Context) {
if req.Role == "" {
req.Role = "user"
}
ctx := c.Request.Context()
userInfo, err := h.service.CreateUser(req.Username, req.Password, req.Role)
userInfo, err := h.service.CreateUser(ctx, req.Username, req.Password, req.Role)
if err != nil {
common.ErrorWithCode(c, common.CodeServerError, err.Error())
return

8
internal/admin/service.go
View File

@@ -170,7 +170,7 @@ func (s *Service) ListUsers(pageIndex, pageSize int, name, status, sort, orderBy
// Returns:
// - map[string]interface{}: user information without password
// - error: error message
func (s *Service) CreateUser(username, password, role string) (map[string]interface{}, error) {
func (s *Service) CreateUser(ctx context.Context, username, password, role string) (map[string]interface{}, error) {
emailRegex := regexp.MustCompile(`^[\w\._-]+@([\w_-]+\.)+[\w-]{2,}$`)
if !emailRegex.MatchString(username) {
return nil, fmt.Errorf("invalid email address: %s", username)
@@ -285,7 +285,7 @@ func (s *Service) CreateUser(username, password, role string) (map[string]interf
}
// 4. Create tenant LLM configurations
tenantLLMs, err := s.getInitTenantLLM(userID)
tenantLLMs, err := s.getInitTenantLLM(ctx, userID)
if err != nil {
common.Warn("failed to get init tenant LLM configs", zap.Error(err))
// Continue without LLM configs - not a critical error
@@ -333,7 +333,7 @@ func (s *Service) CreateUser(username, password, role string) (map[string]interf
// getInitTenantLLM gets initial tenant LLM configurations
// This matches Python's get_init_tenant_llm function
func (s *Service) getInitTenantLLM(userID string) ([]*entity.TenantLLM, error) {
func (s *Service) getInitTenantLLM(ctx context.Context, userID string) ([]*entity.TenantLLM, error) {
cfg := server.GetConfig()
if cfg == nil {
return nil, fmt.Errorf("config not initialized")
@@ -366,7 +366,7 @@ func (s *Service) getInitTenantLLM(userID string) ([]*entity.TenantLLM, error) {
// Get LLMs for each unique factory
for _, factoryConfig := range uniqueFactories {
models, err := s.llmDAO.GetByFactory(factoryConfig.Factory)
models, err := s.llmDAO.GetByFactory(ctx, dao.DB, factoryConfig.Factory)
if err != nil {
common.Warn("failed to get LLMs for factory", zap.String("factory", factoryConfig.Factory), zap.Error(err))
continue

25
internal/dao/langfuse.go
View File

@@ -17,6 +17,7 @@
package dao
import (
"context"
"errors"
"gorm.io/gorm"
@@ -36,9 +37,9 @@ func NewLangfuse() *LangfuseDAO {
// GetByTenantID returns the Langfuse credentials row for a tenant.
// It returns (nil, nil) when no row exists, mirroring the Python
// TenantLangfuseService.filter_by_tenant behaviour (DoesNotExist -> None).
func (dao *LangfuseDAO) GetByTenantID(tenantID string) (*entity.TenantLangfuse, error) {
func (dao *LangfuseDAO) GetByTenantID(ctx context.Context, db *gorm.DB, tenantID string) (*entity.TenantLangfuse, error) {
var row entity.TenantLangfuse
err := DB.Where("tenant_id = ?", tenantID).First(&row).Error
err := db.WithContext(ctx).Where("tenant_id = ?", tenantID).First(&row).Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, nil
@@ -49,13 +50,13 @@ func (dao *LangfuseDAO) GetByTenantID(tenantID string) (*entity.TenantLangfuse,
}
// Create inserts a new Langfuse credentials row (mirrors save).
func (dao *LangfuseDAO) Create(row *entity.TenantLangfuse) error {
return DB.Create(row).Error
func (dao *LangfuseDAO) Create(ctx context.Context, db *gorm.DB, row *entity.TenantLangfuse) error {
return db.WithContext(ctx).Create(row).Error
}
// UpdateByTenantID updates the Langfuse credentials row for a tenant
func (dao *LangfuseDAO) UpdateByTenantID(tenantID string, updates map[string]any) error {
res := DB.Model(&entity.TenantLangfuse{}).Where("tenant_id = ?", tenantID).Updates(updates)
func (dao *LangfuseDAO) UpdateByTenantID(ctx context.Context, db *gorm.DB, tenantID string, updates map[string]any) error {
res := db.WithContext(ctx).Model(&entity.TenantLangfuse{}).Where("tenant_id = ?", tenantID).Updates(updates)
if res.Error != nil {
return res.Error
}
@@ -67,8 +68,8 @@ func (dao *LangfuseDAO) UpdateByTenantID(tenantID string, updates map[string]any
// DeleteByTenantID deletes the Langfuse credentials row for a tenant
// (mirrors delete_model / delete_ty_tenant_id).
func (dao *LangfuseDAO) DeleteByTenantID(tenantID string) error {
res := DB.Where("tenant_id = ?", tenantID).Delete(&entity.TenantLangfuse{})
func (dao *LangfuseDAO) DeleteByTenantID(ctx context.Context, db *gorm.DB, tenantID string) error {
res := db.WithContext(ctx).Where("tenant_id = ?", tenantID).Delete(&entity.TenantLangfuse{})
if res.Error != nil {
return res.Error
}
@@ -78,8 +79,8 @@ func (dao *LangfuseDAO) DeleteByTenantID(tenantID string) error {
return nil
}
func (dao *LangfuseDAO) SaveByTenantID(row *entity.TenantLangfuse) error {
return DB.Clauses(clause.OnConflict{
func (dao *LangfuseDAO) SaveByTenantID(ctx context.Context, db *gorm.DB, row *entity.TenantLangfuse) error {
return db.WithContext(ctx).Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "tenant_id"}},
DoUpdates: clause.Assignments(map[string]any{
"secret_key": row.SecretKey,
@@ -89,8 +90,8 @@ func (dao *LangfuseDAO) SaveByTenantID(row *entity.TenantLangfuse) error {
}).Create(row).Error
}
func (dao *LangfuseDAO) DeleteExistingByTenantID(tenantID string) error {
return DB.Transaction(func(tx *gorm.DB) error {
func (dao *LangfuseDAO) DeleteExistingByTenantID(ctx context.Context, db *gorm.DB, tenantID string) error {
return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
var row entity.TenantLangfuse
err := tx.Where("tenant_id = ?", tenantID).First(&row).Error
if err != nil {

16
internal/dao/langfuse_test.go
View File

@@ -45,8 +45,9 @@ func TestLangfuseDAO_GetByTenantID_NotFound(t *testing.T) {
db := setupLangfuseTestDB(t)
pushDB(t, db)
dao := NewLangfuse()
ctx := t.Context()
row, err := dao.GetByTenantID("missing")
row, err := dao.GetByTenantID(ctx, db, "missing")
if err != nil {
t.Fatalf("expected nil error for missing row, got %v", err)
}
@@ -59,6 +60,7 @@ func TestLangfuseDAO_CRUD(t *testing.T) {
db := setupLangfuseTestDB(t)
pushDB(t, db)
dao := NewLangfuse()
ctx := t.Context()
// 1. Create
row := &entity.TenantLangfuse{
@@ -67,12 +69,12 @@ func TestLangfuseDAO_CRUD(t *testing.T) {
PublicKey: "pk-1",
Host: "https://cloud.langfuse.com",
}
if err := dao.Create(row); err != nil {
if err := dao.Create(ctx, db, row); err != nil {
t.Fatalf("failed to create: %v", err)
}
// 2. GetByTenantID
got, err := dao.GetByTenantID("tenant-1")
got, err := dao.GetByTenantID(ctx, db, "tenant-1")
if err != nil {
t.Fatalf("failed to get: %v", err)
}
@@ -93,10 +95,10 @@ func TestLangfuseDAO_CRUD(t *testing.T) {
"public_key": "pk-2",
"host": "https://eu.langfuse.com",
}
if err := dao.UpdateByTenantID("tenant-1", updates); err != nil {
if err = dao.UpdateByTenantID(ctx, db, "tenant-1", updates); err != nil {
t.Fatalf("failed to update: %v", err)
}
got, err = dao.GetByTenantID("tenant-1")
got, err = dao.GetByTenantID(ctx, db, "tenant-1")
if err != nil {
t.Fatalf("failed to get after update: %v", err)
}
@@ -105,10 +107,10 @@ func TestLangfuseDAO_CRUD(t *testing.T) {
}
// 4. DeleteByTenantID
if err := dao.DeleteByTenantID("tenant-1"); err != nil {
if err = dao.DeleteByTenantID(ctx, db, "tenant-1"); err != nil {
t.Fatalf("failed to delete: %v", err)
}
got, err = dao.GetByTenantID("tenant-1")
got, err = dao.GetByTenantID(ctx, db, "tenant-1")
if err != nil {
t.Fatalf("expected nil error after delete, got %v", err)
}

11
internal/dao/license.go
View File

@@ -17,8 +17,11 @@
package dao
import (
"context"
"ragflow/internal/entity"
"time"
"gorm.io/gorm"
)
// LicenseDAO license data access object
@@ -30,19 +33,19 @@ func NewLicenseDAO() *LicenseDAO {
}
// Create creates a new license record
func (dao *LicenseDAO) Create(licenseID, licenseStr string) error {
func (dao *LicenseDAO) Create(ctx context.Context, db *gorm.DB, licenseID, licenseStr string) error {
license := entity.License{
ID: licenseID,
License: licenseStr,
CreatedAt: time.Now(),
}
return DB.Create(license).Error
return db.WithContext(ctx).Create(license).Error
}
// GetLatest gets the latest license record by creation time
func (dao *LicenseDAO) GetLatest() (*entity.License, error) {
func (dao *LicenseDAO) GetLatest(ctx context.Context, db *gorm.DB) (*entity.License, error) {
var license entity.License
err := DB.Order("created_at DESC").First(&license).Error
err := db.WithContext(ctx).Order("created_at DESC").First(&license).Error
if err != nil {
return nil, err
}

27
internal/dao/llm.go
View File

@@ -17,7 +17,10 @@
package dao
import (
"context"
"ragflow/internal/entity"
"gorm.io/gorm"
)
// LLMDAO LLM data access object
@@ -29,9 +32,9 @@ func NewLLMDAO() *LLMDAO {
}
// GetAll gets all LLMs
func (dao *LLMDAO) GetAll() ([]*entity.LLM, error) {
func (dao *LLMDAO) GetAll(ctx context.Context, db *gorm.DB) ([]*entity.LLM, error) {
var llms []*entity.LLM
err := DB.Find(&llms).Error
err := db.WithContext(ctx).Find(&llms).Error
if err != nil {
return nil, err
}
@@ -39,9 +42,9 @@ func (dao *LLMDAO) GetAll() ([]*entity.LLM, error) {
}
// GetAllValid gets all valid LLMs
func (dao *LLMDAO) GetAllValid() ([]*entity.LLM, error) {
func (dao *LLMDAO) GetAllValid(ctx context.Context, db *gorm.DB) ([]*entity.LLM, error) {
var llms []*entity.LLM
err := DB.Where("status = ?", "1").Find(&llms).Error
err := db.WithContext(ctx).Where("status = ?", "1").Find(&llms).Error
if err != nil {
return nil, err
}
@@ -49,9 +52,9 @@ func (dao *LLMDAO) GetAllValid() ([]*entity.LLM, error) {
}
// GetByFactory gets LLMs by factory
func (dao *LLMDAO) GetByFactory(factory string) ([]*entity.LLM, error) {
func (dao *LLMDAO) GetByFactory(ctx context.Context, db *gorm.DB, factory string) ([]*entity.LLM, error) {
var llms []*entity.LLM
err := DB.Where("fid = ?", factory).Find(&llms).Error
err := db.WithContext(ctx).Where("fid = ?", factory).Find(&llms).Error
if err != nil {
return nil, err
}
@@ -59,9 +62,9 @@ func (dao *LLMDAO) GetByFactory(factory string) ([]*entity.LLM, error) {
}
// GetByFactoryAndName gets LLM by factory and name
func (dao *LLMDAO) GetByFactoryAndName(factory, name string) (*entity.LLM, error) {
func (dao *LLMDAO) GetByFactoryAndName(ctx context.Context, db *gorm.DB, factory, name string) (*entity.LLM, error) {
var llm entity.LLM
err := DB.Where("fid = ? AND llm_name = ?", factory, name).First(&llm).Error
err := db.WithContext(ctx).Where("fid = ? AND llm_name = ?", factory, name).First(&llm).Error
if err != nil {
return nil, err
}
@@ -77,9 +80,9 @@ func NewLLMFactoryDAO() *LLMFactoryDAO {
}
// GetAllValid gets all valid LLM factories
func (dao *LLMFactoryDAO) GetAllValid() ([]*entity.LLMFactories, error) {
func (dao *LLMFactoryDAO) GetAllValid(ctx context.Context, db *gorm.DB) ([]*entity.LLMFactories, error) {
var factories []*entity.LLMFactories
err := DB.Where("status = ?", "1").Find(&factories).Error
err := db.WithContext(ctx).Where("status = ?", "1").Find(&factories).Error
if err != nil {
return nil, err
}
@@ -87,9 +90,9 @@ func (dao *LLMFactoryDAO) GetAllValid() ([]*entity.LLMFactories, error) {
}
// GetByName gets LLM factory by name
func (dao *LLMFactoryDAO) GetByName(name string) (*entity.LLMFactories, error) {
func (dao *LLMFactoryDAO) GetByName(ctx context.Context, db *gorm.DB, name string) (*entity.LLMFactories, error) {
var factory entity.LLMFactories
err := DB.Where("name = ?", name).First(&factory).Error
err := db.WithContext(ctx).Where("name = ?", name).First(&factory).Error
if err != nil {
return nil, err
}

17
internal/handler/langfuse.go
View File

@@ -17,6 +17,8 @@
package handler
import (
"context"
"github.com/gin-gonic/gin"
"ragflow/internal/common"
@@ -27,9 +29,9 @@ import (
// LangfuseService is the behaviour the handler depends on (interface enables
// mocking in tests).
type LangfuseService interface {
SetAPIKey(tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error)
GetAPIKey(tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error)
DeleteAPIKey(tenantID string) (bool, common.ErrorCode, string, error)
SetAPIKey(ctx context.Context, tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error)
GetAPIKey(ctx context.Context, tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error)
DeleteAPIKey(ctx context.Context, tenantID string) (bool, common.ErrorCode, string, error)
}
// LangfuseHandler handles /langfuse/api-key HTTP requests.
@@ -68,8 +70,9 @@ func (h *LangfuseHandler) SetAPIKey(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeDataError, nil, "Invalid request: "+err.Error())
return
}
ctx := c.Request.Context()
row, code, err := h.langfuseService.SetAPIKey(user.ID, req.SecretKey, req.PublicKey, req.Host)
row, code, err := h.langfuseService.SetAPIKey(ctx, user.ID, req.SecretKey, req.PublicKey, req.Host)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -91,8 +94,9 @@ func (h *LangfuseHandler) GetAPIKey(c *gin.Context) {
common.ErrorWithCode(c, errorCode, errorMessage)
return
}
ctx := c.Request.Context()
data, code, message, err := h.langfuseService.GetAPIKey(user.ID)
data, code, message, err := h.langfuseService.GetAPIKey(ctx, user.ID)
if err != nil {
common.ResponseWithCodeData(c, code, nil, message)
return
@@ -107,8 +111,9 @@ func (h *LangfuseHandler) DeleteAPIKey(c *gin.Context) {
common.ErrorWithCode(c, errorCode, errorMessage)
return
}
ctx := c.Request.Context()
ok, code, message, err := h.langfuseService.DeleteAPIKey(user.ID)
ok, code, message, err := h.langfuseService.DeleteAPIKey(ctx, user.ID)
if err != nil {
common.ResponseWithCodeData(c, code, nil, message)
return

35
internal/handler/langfuse_test.go
View File

@@ -17,6 +17,7 @@
package handler
import (
"context"
"encoding/json"
"errors"
"net/http"
@@ -31,30 +32,30 @@ import (
)
type fakeLangfuseService struct {
setFn func(tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error)
getFn func(tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error)
deleteFn func(tenantID string) (bool, common.ErrorCode, string, error)
setFn func(ctx context.Context, tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error)
getFn func(ctx context.Context, tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error)
deleteFn func(ctx context.Context, tenantID string) (bool, common.ErrorCode, string, error)
}
func (f fakeLangfuseService) SetAPIKey(tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
func (f fakeLangfuseService) SetAPIKey(ctx context.Context, tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
if f.setFn == nil {
return nil, common.CodeServerError, errors.New("unexpected SetAPIKey call")
}
return f.setFn(tenantID, secretKey, publicKey, host)
return f.setFn(ctx, tenantID, secretKey, publicKey, host)
}
func (f fakeLangfuseService) GetAPIKey(tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
func (f fakeLangfuseService) GetAPIKey(ctx context.Context, tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
if f.getFn == nil {
return nil, common.CodeServerError, "", errors.New("unexpected GetAPIKey call")
}
return f.getFn(tenantID)
return f.getFn(ctx, tenantID)
}
func (f fakeLangfuseService) DeleteAPIKey(tenantID string) (bool, common.ErrorCode, string, error) {
func (f fakeLangfuseService) DeleteAPIKey(ctx context.Context, tenantID string) (bool, common.ErrorCode, string, error) {
if f.deleteFn == nil {
return false, common.CodeServerError, "", errors.New("unexpected DeleteAPIKey call")
}
return f.deleteFn(tenantID)
return f.deleteFn(ctx, tenantID)
}
func serveLangfuse(method, target, body string, h func(c *gin.Context)) *httptest.ResponseRecorder {
@@ -89,7 +90,7 @@ func decode(t *testing.T, resp *httptest.ResponseRecorder) map[string]interface{
func TestLangfuseHandler_SetAPIKey_Success(t *testing.T) {
var gotTenant, gotSecret, gotPublic, gotHost string
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
setFn: func(tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
setFn: func(ctx context.Context, tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
gotTenant, gotSecret, gotPublic, gotHost = tenantID, secretKey, publicKey, host
return &entity.TenantLangfuse{TenantID: tenantID, SecretKey: secretKey, PublicKey: publicKey, Host: host}, common.CodeSuccess, nil
},
@@ -116,7 +117,7 @@ func TestLangfuseHandler_SetAPIKey_Success(t *testing.T) {
func TestLangfuseHandler_SetAPIKey_ServiceError(t *testing.T) {
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
setFn: func(tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
setFn: func(ctx context.Context, tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
return nil, common.CodeDataError, errors.New("Invalid Langfuse keys")
},
}}
@@ -136,7 +137,7 @@ func TestLangfuseHandler_SetAPIKey_ServiceError(t *testing.T) {
func TestLangfuseHandler_SetAPIKey_BindFailureStopsEarly(t *testing.T) {
called := false
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
setFn: func(tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
setFn: func(ctx context.Context, tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
called = true
return nil, common.CodeSuccess, nil
},
@@ -155,7 +156,7 @@ func TestLangfuseHandler_SetAPIKey_BindFailureStopsEarly(t *testing.T) {
func TestLangfuseHandler_GetAPIKey_Success(t *testing.T) {
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
getFn: func(tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
getFn: func(ctx context.Context, tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
return &entity.LangfuseInfoResponse{
TenantID: tenantID, Host: "host", SecretKey: "sk", PublicKey: "pk",
ProjectID: "proj-1", ProjectName: "My Project",
@@ -180,7 +181,7 @@ func TestLangfuseHandler_GetAPIKey_Success(t *testing.T) {
func TestLangfuseHandler_GetAPIKey_NoRecord(t *testing.T) {
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
getFn: func(tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
getFn: func(ctx context.Context, tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
return nil, common.CodeSuccess, "Have not record any Langfuse keys.", nil
},
}}
@@ -201,7 +202,7 @@ func TestLangfuseHandler_GetAPIKey_NoRecord(t *testing.T) {
func TestLangfuseHandler_GetAPIKey_Unauthorized(t *testing.T) {
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
getFn: func(tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
getFn: func(ctx context.Context, tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
return nil, common.CodeDataError, "Invalid Langfuse keys loaded", errors.New("unauthorized")
},
}}
@@ -220,7 +221,7 @@ func TestLangfuseHandler_GetAPIKey_Unauthorized(t *testing.T) {
func TestLangfuseHandler_DeleteAPIKey_Success(t *testing.T) {
var gotTenant string
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
deleteFn: func(tenantID string) (bool, common.ErrorCode, string, error) {
deleteFn: func(ctx context.Context, tenantID string) (bool, common.ErrorCode, string, error) {
gotTenant = tenantID
return true, common.CodeSuccess, "", nil
},
@@ -242,7 +243,7 @@ func TestLangfuseHandler_DeleteAPIKey_Success(t *testing.T) {
func TestLangfuseHandler_DeleteAPIKey_NoRecord(t *testing.T) {
h := &LangfuseHandler{langfuseService: fakeLangfuseService{
deleteFn: func(tenantID string) (bool, common.ErrorCode, string, error) {
deleteFn: func(ctx context.Context, tenantID string) (bool, common.ErrorCode, string, error) {
return false, common.CodeSuccess, "Have not record any Langfuse keys.", nil
},
}}

10
internal/handler/llm.go
View File

@@ -67,8 +67,9 @@ func (h *LLMHandler) GetMyLLMs(c *gin.Context) {
tenantID := user.ID
includeDetailsStr := c.DefaultQuery("include_details", "false")
includeDetails := includeDetailsStr == "true"
ctx := c.Request.Context()
llms, err := h.llmService.GetMyLLMs(tenantID, includeDetails)
llms, err := h.llmService.GetMyLLMs(ctx, tenantID, includeDetails)
if err != nil {
common.ResponseWithCodeData(c, common.CodeExceptionError, false, err.Error())
return
@@ -100,8 +101,10 @@ func (h *LLMHandler) SetAPIKey(c *gin.Context) {
return
}
ctx := c.Request.Context()
tenantID := user.ID
result, err := h.llmService.SetAPIKey(tenantID, &req)
result, err := h.llmService.SetAPIKey(ctx, tenantID, &req)
if err != nil {
common.ResponseWithCodeData(c, common.CodeDataError, false, err.Error())
return
@@ -133,10 +136,11 @@ func (h *LLMHandler) ListApp(c *gin.Context) {
}
tenantID := user.ID
ctx := c.Request.Context()
modelType := c.Query("model_type")
llms, err := h.llmService.ListLLMs(tenantID, modelType)
llms, err := h.llmService.ListLLMs(ctx, tenantID, modelType)
if err != nil {
common.ResponseWithCodeData(c, common.CodeExceptionError, false, err.Error())
return

2
internal/handler/user.go
View File

@@ -63,7 +63,7 @@ func (h *UserHandler) Register(c *gin.Context) {
return
}
user, code, err := h.userService.Register(&req)
user, code, err := h.userService.Register(ctx, &req)
if err != nil {
var data interface{} = false
if code == common.CodeExceptionError {

20
internal/service/langfuse_service.go
View File

@@ -69,17 +69,17 @@ func NewLangfuseService() *LangfuseService {
// SetAPIKey validates and stores (insert or update) the Langfuse credentials
// for a tenant.
func (s *LangfuseService) SetAPIKey(tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
func (s *LangfuseService) SetAPIKey(ctx context.Context, tenantID, secretKey, publicKey, host string) (*entity.TenantLangfuse, common.ErrorCode, error) {
if secretKey == "" || publicKey == "" || host == "" {
return nil, common.CodeDataError, errors.New("Missing required fields")
return nil, common.CodeDataError, errors.New("missing required fields")
}
ok, err := s.verifier.AuthCheck(context.Background(), host, publicKey, secretKey)
ok, err := s.verifier.AuthCheck(ctx, host, publicKey, secretKey)
if err != nil {
return nil, common.CodeServerError, err
}
if !ok {
return nil, common.CodeDataError, errors.New("Invalid Langfuse keys")
return nil, common.CodeDataError, errors.New("invalid Langfuse keys")
}
row := &entity.TenantLangfuse{
@@ -89,7 +89,7 @@ func (s *LangfuseService) SetAPIKey(tenantID, secretKey, publicKey, host string)
Host: host,
}
if err := s.langfuseDAO.SaveByTenantID(row); err != nil {
if err = s.langfuseDAO.SaveByTenantID(ctx, dao.DB, row); err != nil {
return nil, common.CodeServerError, err
}
@@ -98,8 +98,8 @@ func (s *LangfuseService) SetAPIKey(tenantID, secretKey, publicKey, host string)
// GetAPIKey returns the stored credentials enriched with the Langfuse project
// id/name.
func (s *LangfuseService) GetAPIKey(tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
row, err := s.langfuseDAO.GetByTenantID(tenantID)
func (s *LangfuseService) GetAPIKey(ctx context.Context, tenantID string) (*entity.LangfuseInfoResponse, common.ErrorCode, string, error) {
row, err := s.langfuseDAO.GetByTenantID(ctx, dao.DB, tenantID)
if err != nil {
return nil, common.CodeServerError, "", err
}
@@ -107,7 +107,7 @@ func (s *LangfuseService) GetAPIKey(tenantID string) (*entity.LangfuseInfoRespon
return nil, common.CodeSuccess, "Have not record any Langfuse keys.", nil
}
projectID, projectName, err := s.verifier.GetProject(context.Background(), row.Host, row.PublicKey, row.SecretKey)
projectID, projectName, err := s.verifier.GetProject(ctx, row.Host, row.PublicKey, row.SecretKey)
if err != nil {
if errors.Is(err, ErrLangfuseUnauthorized) {
return nil, common.CodeDataError, "Invalid Langfuse keys loaded", err
@@ -130,8 +130,8 @@ func (s *LangfuseService) GetAPIKey(tenantID string) (*entity.LangfuseInfoRespon
}
// DeleteAPIKey removes the stored credentials for a tenant.
func (s *LangfuseService) DeleteAPIKey(tenantID string) (bool, common.ErrorCode, string, error) {
if err := s.langfuseDAO.DeleteExistingByTenantID(tenantID); err != nil {
func (s *LangfuseService) DeleteAPIKey(ctx context.Context, tenantID string) (bool, common.ErrorCode, string, error) {
if err := s.langfuseDAO.DeleteExistingByTenantID(ctx, dao.DB, tenantID); err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return false, common.CodeSuccess, "Have not record any Langfuse keys.", nil
}

59
internal/service/langfuse_service_test.go
View File

@@ -72,9 +72,10 @@ func newLangfuseServiceForTest(v langfuseVerifier) *LangfuseService {
func TestLangfuseService_SetAPIKey_MissingFields(t *testing.T) {
setupLangfuseServiceTestDB(t)
svc := newLangfuseServiceForTest(stubLangfuseVerifier{authOK: true})
ctx := t.Context()
_, code, err := svc.SetAPIKey("tenant-1", "", "pk", "host")
if code != common.CodeDataError || err == nil || err.Error() != "Missing required fields" {
_, code, err := svc.SetAPIKey(ctx, "tenant-1", "", "pk", "host")
if code != common.CodeDataError || err == nil || err.Error() != "missing required fields" {
t.Fatalf("expected Missing required fields/CodeDataError, got code=%d err=%v", code, err)
}
}
@@ -82,9 +83,10 @@ func TestLangfuseService_SetAPIKey_MissingFields(t *testing.T) {
func TestLangfuseService_SetAPIKey_InvalidKeys(t *testing.T) {
setupLangfuseServiceTestDB(t)
svc := newLangfuseServiceForTest(stubLangfuseVerifier{authOK: false})
ctx := t.Context()
_, code, err := svc.SetAPIKey("tenant-1", "sk", "pk", "host")
if code != common.CodeDataError || err == nil || err.Error() != "Invalid Langfuse keys" {
_, code, err := svc.SetAPIKey(ctx, "tenant-1", "sk", "pk", "host")
if code != common.CodeDataError || err == nil || err.Error() != "invalid Langfuse keys" {
t.Fatalf("expected Invalid Langfuse keys/CodeDataError, got code=%d err=%v", code, err)
}
}
@@ -92,8 +94,9 @@ func TestLangfuseService_SetAPIKey_InvalidKeys(t *testing.T) {
func TestLangfuseService_SetAPIKey_VerifierError(t *testing.T) {
setupLangfuseServiceTestDB(t)
svc := newLangfuseServiceForTest(stubLangfuseVerifier{authErr: errors.New("network down")})
ctx := t.Context()
_, code, err := svc.SetAPIKey("tenant-1", "sk", "pk", "host")
_, code, err := svc.SetAPIKey(ctx, "tenant-1", "sk", "pk", "host")
if code != common.CodeServerError || err == nil || err.Error() != "network down" {
t.Fatalf("expected verifier error/CodeServerError, got code=%d err=%v", code, err)
}
@@ -102,9 +105,10 @@ func TestLangfuseService_SetAPIKey_VerifierError(t *testing.T) {
func TestLangfuseService_SetAPIKey_CreateThenUpdate(t *testing.T) {
db := setupLangfuseServiceTestDB(t)
svc := newLangfuseServiceForTest(stubLangfuseVerifier{authOK: true})
ctx := t.Context()
// Create
row, code, err := svc.SetAPIKey("tenant-1", "sk-1", "pk-1", "https://a.langfuse.com")
row, code, err := svc.SetAPIKey(ctx, "tenant-1", "sk-1", "pk-1", "https://a.langfuse.com")
if err != nil || code != common.CodeSuccess {
t.Fatalf("create failed: code=%d err=%v", code, err)
}
@@ -119,7 +123,7 @@ func TestLangfuseService_SetAPIKey_CreateThenUpdate(t *testing.T) {
}
// Update (same tenant) should not create a second row
_, code, err = svc.SetAPIKey("tenant-1", "sk-2", "pk-2", "https://b.langfuse.com")
_, code, err = svc.SetAPIKey(ctx, "tenant-1", "sk-2", "pk-2", "https://b.langfuse.com")
if err != nil || code != common.CodeSuccess {
t.Fatalf("update failed: code=%d err=%v", code, err)
}
@@ -127,7 +131,7 @@ func TestLangfuseService_SetAPIKey_CreateThenUpdate(t *testing.T) {
if count != 1 {
t.Fatalf("expected still 1 row after update, got %d", count)
}
stored, _ := dao.NewLangfuse().GetByTenantID("tenant-1")
stored, _ := dao.NewLangfuse().GetByTenantID(ctx, db, "tenant-1")
if stored == nil || stored.SecretKey != "sk-2" || stored.Host != "https://b.langfuse.com" {
t.Fatalf("update not persisted: %+v", stored)
}
@@ -136,8 +140,9 @@ func TestLangfuseService_SetAPIKey_CreateThenUpdate(t *testing.T) {
func TestLangfuseService_GetAPIKey_NoRecord(t *testing.T) {
setupLangfuseServiceTestDB(t)
svc := newLangfuseServiceForTest(stubLangfuseVerifier{})
ctx := t.Context()
data, code, message, err := svc.GetAPIKey("tenant-1")
data, code, message, err := svc.GetAPIKey(ctx, "tenant-1")
if err != nil || code != common.CodeSuccess || data != nil {
t.Fatalf("unexpected: code=%d data=%v err=%v", code, data, err)
}
@@ -147,13 +152,14 @@ func TestLangfuseService_GetAPIKey_NoRecord(t *testing.T) {
}
func TestLangfuseService_GetAPIKey_Unauthorized(t *testing.T) {
setupLangfuseServiceTestDB(t)
if err := dao.NewLangfuse().Create(&entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
db := setupLangfuseServiceTestDB(t)
ctx := t.Context()
if err := dao.NewLangfuse().Create(ctx, db, &entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
t.Fatalf("seed failed: %v", err)
}
svc := newLangfuseServiceForTest(stubLangfuseVerifier{projErr: ErrLangfuseUnauthorized})
data, code, message, err := svc.GetAPIKey("tenant-1")
data, code, message, err := svc.GetAPIKey(ctx, "tenant-1")
if data != nil || code != common.CodeDataError || err == nil {
t.Fatalf("unexpected: code=%d data=%v err=%v", code, data, err)
}
@@ -163,13 +169,14 @@ func TestLangfuseService_GetAPIKey_Unauthorized(t *testing.T) {
}
func TestLangfuseService_GetAPIKey_ApiError(t *testing.T) {
setupLangfuseServiceTestDB(t)
if err := dao.NewLangfuse().Create(&entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
db := setupLangfuseServiceTestDB(t)
ctx := t.Context()
if err := dao.NewLangfuse().Create(ctx, db, &entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
t.Fatalf("seed failed: %v", err)
}
svc := newLangfuseServiceForTest(stubLangfuseVerifier{projErr: &LangfuseAPIError{StatusCode: 500, Body: "boom"}})
data, code, message, err := svc.GetAPIKey("tenant-1")
data, code, message, err := svc.GetAPIKey(ctx, "tenant-1")
if data != nil || code != common.CodeSuccess || err != nil {
t.Fatalf("unexpected: code=%d data=%v err=%v", code, data, err)
}
@@ -179,26 +186,28 @@ func TestLangfuseService_GetAPIKey_ApiError(t *testing.T) {
}
func TestLangfuseService_GetAPIKey_NonAPIError(t *testing.T) {
setupLangfuseServiceTestDB(t)
if err := dao.NewLangfuse().Create(&entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
db := setupLangfuseServiceTestDB(t)
ctx := t.Context()
if err := dao.NewLangfuse().Create(ctx, db, &entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
t.Fatalf("seed failed: %v", err)
}
svc := newLangfuseServiceForTest(stubLangfuseVerifier{projErr: errors.New("json parse failed")})
data, code, message, err := svc.GetAPIKey("tenant-1")
data, code, message, err := svc.GetAPIKey(ctx, "tenant-1")
if data != nil || code != common.CodeServerError || message != "" || err == nil {
t.Fatalf("unexpected: code=%d message=%q data=%v err=%v", code, message, data, err)
}
}
func TestLangfuseService_GetAPIKey_Success(t *testing.T) {
setupLangfuseServiceTestDB(t)
if err := dao.NewLangfuse().Create(&entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "https://a.langfuse.com"}); err != nil {
db := setupLangfuseServiceTestDB(t)
ctx := t.Context()
if err := dao.NewLangfuse().Create(ctx, db, &entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "https://a.langfuse.com"}); err != nil {
t.Fatalf("seed failed: %v", err)
}
svc := newLangfuseServiceForTest(stubLangfuseVerifier{projID: "proj-1", projName: "My Project"})
data, code, message, err := svc.GetAPIKey("tenant-1")
data, code, message, err := svc.GetAPIKey(ctx, "tenant-1")
if err != nil || code != common.CodeSuccess || message != "success" {
t.Fatalf("unexpected: code=%d message=%q err=%v", code, message, err)
}
@@ -215,8 +224,9 @@ func TestLangfuseService_GetAPIKey_Success(t *testing.T) {
func TestLangfuseService_DeleteAPIKey_NoRecord(t *testing.T) {
setupLangfuseServiceTestDB(t)
svc := newLangfuseServiceForTest(stubLangfuseVerifier{})
ctx := t.Context()
ok, code, message, err := svc.DeleteAPIKey("tenant-1")
ok, code, message, err := svc.DeleteAPIKey(ctx, "tenant-1")
if ok || code != common.CodeSuccess || err != nil {
t.Fatalf("unexpected: ok=%v code=%d err=%v", ok, code, err)
}
@@ -227,12 +237,13 @@ func TestLangfuseService_DeleteAPIKey_NoRecord(t *testing.T) {
func TestLangfuseService_DeleteAPIKey_Success(t *testing.T) {
db := setupLangfuseServiceTestDB(t)
if err := dao.NewLangfuse().Create(&entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
ctx := t.Context()
if err := dao.NewLangfuse().Create(ctx, db, &entity.TenantLangfuse{TenantID: "tenant-1", SecretKey: "sk", PublicKey: "pk", Host: "host"}); err != nil {
t.Fatalf("seed failed: %v", err)
}
svc := newLangfuseServiceForTest(stubLangfuseVerifier{})
ok, code, message, err := svc.DeleteAPIKey("tenant-1")
ok, code, message, err := svc.DeleteAPIKey(ctx, "tenant-1")
if !ok || code != common.CodeSuccess || message != "" || err != nil {
t.Fatalf("unexpected: ok=%v code=%d message=%q err=%v", ok, code, message, err)
}

13
internal/service/llm.go
View File

@@ -17,6 +17,7 @@
package service
import (
"context"
"fmt"
"ragflow/internal/entity"
"strconv"
@@ -57,7 +58,7 @@ type MyLLMFactory struct {
}
// GetMyLLMs get my LLMs for a tenant
func (s *LLMService) GetMyLLMs(tenantID string, includeDetails bool) (map[string]MyLLMFactory, error) {
func (s *LLMService) GetMyLLMs(ctx context.Context, tenantID string, includeDetails bool) (map[string]MyLLMFactory, error) {
result := make(map[string]MyLLMFactory)
if includeDetails {
@@ -67,7 +68,7 @@ func (s *LLMService) GetMyLLMs(tenantID string, includeDetails bool) (map[string
}
factoryDAO := dao.NewLLMFactoryDAO()
factories, err := factoryDAO.GetAllValid()
factories, err := factoryDAO.GetAllValid(ctx, dao.DB)
if err != nil {
return nil, err
}
@@ -159,7 +160,7 @@ type LLMListItem struct {
type ListLLMsResponse map[string][]LLMListItem
// ListLLMs lists LLMs for a tenant with availability info
func (s *LLMService) ListLLMs(tenantID string, modelType string) (ListLLMsResponse, error) {
func (s *LLMService) ListLLMs(ctx context.Context, tenantID string, modelType string) (ListLLMsResponse, error) {
selfDeployed := map[string]bool{
"FastEmbed": true,
"Ollama": true,
@@ -191,7 +192,7 @@ func (s *LLMService) ListLLMs(tenantID string, modelType string) (ListLLMsRespon
tenantLLMMapping[key] = int64ToString(o.ID)
}
allLLMs, err := s.llmDAO.GetAllValid()
allLLMs, err := s.llmDAO.GetAllValid(ctx, dao.DB)
if err != nil {
return nil, err
}
@@ -337,7 +338,7 @@ type SetAPIKeyResult struct {
}
// SetAPIKey sets API key for a LLM factory
func (s *LLMService) SetAPIKey(tenantID string, req *SetAPIKeyRequest) (*SetAPIKeyResult, error) {
func (s *LLMService) SetAPIKey(ctx context.Context, tenantID string, req *SetAPIKeyRequest) (*SetAPIKeyResult, error) {
factory := req.LLMFactory
baseURL := req.BaseURL
sourceFactory := req.SourceFID
@@ -345,7 +346,7 @@ func (s *LLMService) SetAPIKey(tenantID string, req *SetAPIKeyRequest) (*SetAPIK
sourceFactory = factory
}
sourceLLMs, err := s.llmDAO.GetByFactory(sourceFactory)
sourceLLMs, err := s.llmDAO.GetByFactory(ctx, dao.DB, sourceFactory)
if err != nil || len(sourceLLMs) == 0 {
msg := "No models configured for " + factory + " (source: " + sourceFactory + ")."
if req.Verify {

8
internal/service/user.go
View File

@@ -102,7 +102,7 @@ type UserResponse struct {
}
// Register user registration
func (s *UserService) Register(req *RegisterRequest) (*entity.User, common.ErrorCode, error) {
func (s *UserService) Register(ctx context.Context, req *RegisterRequest) (*entity.User, common.ErrorCode, error) {
cfg := server.GetConfig()
if !cfg.Authentication.RegisterEnabled {
return nil, common.CodeOperatingError, fmt.Errorf("User registration is disabled!")
@@ -226,7 +226,7 @@ func (s *UserService) Register(req *RegisterRequest) (*entity.User, common.Error
Size: 0,
}
tenantLLMs, err := s.getInitTenantLLM(userID)
tenantLLMs, err := s.getInitTenantLLM(ctx, userID)
if err != nil {
return nil, common.CodeServerError, fmt.Errorf("failed to initialize tenant llm: %w", err)
}
@@ -262,7 +262,7 @@ func (s *UserService) Register(req *RegisterRequest) (*entity.User, common.Error
}
// getInitTenantLLM builds the tenant_llm rows created for a new user's default tenant.
func (s *UserService) getInitTenantLLM(userID string) ([]*entity.TenantLLM, error) {
func (s *UserService) getInitTenantLLM(ctx context.Context, userID string) ([]*entity.TenantLLM, error) {
cfg := server.GetConfig()
if cfg == nil {
return nil, fmt.Errorf("config not initialized")
@@ -295,7 +295,7 @@ func (s *UserService) getInitTenantLLM(userID string) ([]*entity.TenantLLM, erro
llmDAO := dao.NewLLMDAO()
tenantLLMs := make([]*entity.TenantLLM, 0)
for _, factoryConfig := range factoryConfigs {
llms, err := llmDAO.GetByFactory(factoryConfig.Factory)
llms, err := llmDAO.GetByFactory(ctx, dao.DB, factoryConfig.Factory)
if err != nil {
return nil, fmt.Errorf("failed to get LLMs for factory %s: %w", factoryConfig.Factory, err)
}