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

Go: add context to lots of interface (#17253)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-22 22:30:57 +08:00
committed by GitHub
parent b3d394954d
commit d19a036cda
221 changed files with 3215 additions and 2337 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
internal/handler/auth.go
View File

@@ -17,6 +17,7 @@
package handler
import (
"context"
"fmt"
"net/http"
"ragflow/internal/common"
@@ -37,10 +38,10 @@ type AuthHandler struct {
// so the test suite can swap in a stub without spinning up the
// full UserService (which requires a live Redis + JWT secret).
type userTokenResolver interface {
GetUserByToken(authorization string) (*entity.User, common.ErrorCode, error)
GetUserByAPIToken(token string) (*entity.User, common.ErrorCode, error)
GetUserByBetaAPIToken(token string) (*entity.User, common.ErrorCode, error)
GetAPITokenByBeta(authorization string) (*entity.APIToken, error)
GetUserByToken(ctx context.Context, authorization string) (*entity.User, common.ErrorCode, error)
GetUserByAPIToken(ctx context.Context, token string) (*entity.User, common.ErrorCode, error)
GetUserByBetaAPIToken(ctx context.Context, token string) (*entity.User, common.ErrorCode, error)
GetAPITokenByBeta(ctx context.Context, authorization string) (*entity.APIToken, error)
}
// NewAuthHandler create auth handler
@@ -69,6 +70,7 @@ func NewAuthHandler() *AuthHandler {
// regular user token must keep working here too.
func (h *AuthHandler) BetaAuthMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
ctx := c.Request.Context()
auth := c.GetHeader("Authorization")
if auth == "" {
if cookie, err := c.Cookie(oauthAuthCookie); err == nil {
@@ -82,13 +84,13 @@ func (h *AuthHandler) BetaAuthMiddleware() gin.HandlerFunc {
return
}
// AUTH_JWT
if u, code, err := h.userService.GetUserByToken(auth); err == nil && code == common.CodeSuccess {
if u, code, err := h.userService.GetUserByToken(ctx, auth); err == nil && code == common.CodeSuccess {
c.Set("user", u)
c.Next()
return
}
// Then try a regular API token (non-beta public bot flow).
if u, code, err := h.userService.GetUserByAPIToken(auth); err == nil && code == common.CodeSuccess {
if u, code, err := h.userService.GetUserByAPIToken(ctx, auth); err == nil && code == common.CodeSuccess {
c.Set("user", u)
c.Set("auth_via_api_token", true)
c.Next()
@@ -101,9 +103,9 @@ func (h *AuthHandler) BetaAuthMiddleware() gin.HandlerFunc {
// Mirrors the python
// `APIToken.query(beta=token).dialog_id` lookup in
// bot_api.py:agent_bot_logs.
if u, code, err := h.userService.GetUserByBetaAPIToken(auth); err == nil && code == common.CodeSuccess {
if u, code, err := h.userService.GetUserByBetaAPIToken(ctx, auth); err == nil && code == common.CodeSuccess {
c.Set("user", u)
if tok, terr := h.userService.GetAPITokenByBeta(auth); terr == nil && tok != nil && tok.DialogID != nil {
if tok, terr := h.userService.GetAPITokenByBeta(ctx, auth); terr == nil && tok != nil && tok.DialogID != nil {
// tok.DialogID is *string (nullable in the schema), but
// downstream handlers (GetAgentbotLogs, GetAgentLogs)
// read "agent_id" with agentID.(string) — they cannot
@@ -126,6 +128,7 @@ func (h *AuthHandler) BetaAuthMiddleware() gin.HandlerFunc {
// Validates that the user is authenticated and is a superuser (admin)
func (h *AuthHandler) AuthMiddleware() gin.HandlerFunc {
return func(c *gin.Context) {
ctx := c.Request.Context()
token := c.GetHeader("Authorization")
if token == "" {
common.ResponseWithHttpCodeData(c, http.StatusUnauthorized, 401, nil, "Missing Authorization header")
@@ -136,9 +139,9 @@ func (h *AuthHandler) AuthMiddleware() gin.HandlerFunc {
authViaAPIToken := false
// Get user by access token
user, code, err := h.userService.GetUserByToken(token)
user, code, err := h.userService.GetUserByToken(ctx, token)
if err != nil {
user, code, err = h.userService.GetUserByAPIToken(token)
user, code, err = h.userService.GetUserByAPIToken(ctx, token)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusUnauthorized, code, nil, "Invalid access token")
c.Abort()

33
internal/handler/bot_test.go
View File

@@ -784,7 +784,7 @@ func TestBotRoutes_RequireAuth(t *testing.T) {
func TestBotMiddleware_NonBearerRegularToken(t *testing.T) {
gin.SetMode(gin.TestMode)
stub := &stubUserTokenResolver{
getUserByTokenFn: func(auth string) (*entity.User, common.ErrorCode, error) {
getUserByTokenFn: func(ctx context.Context, auth string) (*entity.User, common.ErrorCode, error) {
if auth != "raw-access-token-abc" {
t.Errorf("GetUserByToken called with %q, want raw-access-token-abc", auth)
}
@@ -820,36 +820,36 @@ func TestBotMiddleware_NonBearerRegularToken(t *testing.T) {
// return safe defaults (CodeUnauthorized so the middleware
// short-circuits to 401).
type stubUserTokenResolver struct {
getUserByTokenFn func(authorization string) (*entity.User, common.ErrorCode, error)
getUserByAPITokenFn func(token string) (*entity.User, common.ErrorCode, error)
getUserByBetaAPITokenFn func(token string) (*entity.User, common.ErrorCode, error)
getAPITokenByBetaFn func(authorization string) (*entity.APIToken, error)
getUserByTokenFn func(ctx context.Context, authorization string) (*entity.User, common.ErrorCode, error)
getUserByAPITokenFn func(ctx context.Context, token string) (*entity.User, common.ErrorCode, error)
getUserByBetaAPITokenFn func(ctx context.Context, token string) (*entity.User, common.ErrorCode, error)
getAPITokenByBetaFn func(ctx context.Context, authorization string) (*entity.APIToken, error)
}
func (s *stubUserTokenResolver) GetUserByToken(authorization string) (*entity.User, common.ErrorCode, error) {
func (s *stubUserTokenResolver) GetUserByToken(ctx context.Context, authorization string) (*entity.User, common.ErrorCode, error) {
if s.getUserByTokenFn != nil {
return s.getUserByTokenFn(authorization)
return s.getUserByTokenFn(ctx, authorization)
}
return nil, common.CodeUnauthorized, errors.New("not stubbed")
}
func (s *stubUserTokenResolver) GetUserByAPIToken(token string) (*entity.User, common.ErrorCode, error) {
func (s *stubUserTokenResolver) GetUserByAPIToken(ctx context.Context, token string) (*entity.User, common.ErrorCode, error) {
if s.getUserByAPITokenFn != nil {
return s.getUserByAPITokenFn(token)
return s.getUserByAPITokenFn(ctx, token)
}
return nil, common.CodeUnauthorized, errors.New("not stubbed")
}
func (s *stubUserTokenResolver) GetUserByBetaAPIToken(token string) (*entity.User, common.ErrorCode, error) {
func (s *stubUserTokenResolver) GetUserByBetaAPIToken(ctx context.Context, token string) (*entity.User, common.ErrorCode, error) {
if s.getUserByBetaAPITokenFn != nil {
return s.getUserByBetaAPITokenFn(token)
return s.getUserByBetaAPITokenFn(ctx, token)
}
return nil, common.CodeUnauthorized, errors.New("not stubbed")
}
func (s *stubUserTokenResolver) GetAPITokenByBeta(authorization string) (*entity.APIToken, error) {
func (s *stubUserTokenResolver) GetAPITokenByBeta(ctx context.Context, authorization string) (*entity.APIToken, error) {
if s.getAPITokenByBetaFn != nil {
return s.getAPITokenByBetaFn(authorization)
return s.getAPITokenByBetaFn(ctx, authorization)
}
return nil, errors.New("not stubbed")
}
@@ -870,7 +870,7 @@ func (s *stubUserTokenResolver) GetAPITokenByBeta(authorization string) (*entity
func TestBotRoutes_NoRegularAuthRequired(t *testing.T) {
gin.SetMode(gin.TestMode)
stub := &stubUserTokenResolver{
getUserByTokenFn: func(auth string) (*entity.User, common.ErrorCode, error) {
getUserByTokenFn: func(ctx context.Context, auth string) (*entity.User, common.ErrorCode, error) {
return &entity.User{ID: "u-regular"}, common.CodeSuccess, nil
},
}
@@ -955,7 +955,7 @@ func TestBotRoutes_NoRegularAuthRequired(t *testing.T) {
func TestDownloadAttachment_Unauth(t *testing.T) {
gin.SetMode(gin.TestMode)
stub := &stubUserTokenResolver{
getUserByTokenFn: func(auth string) (*entity.User, common.ErrorCode, error) {
getUserByTokenFn: func(ctx context.Context, auth string) (*entity.User, common.ErrorCode, error) {
return nil, common.CodeUnauthorized, errors.New("invalid token")
},
}
@@ -968,13 +968,14 @@ func TestDownloadAttachment_Unauth(t *testing.T) {
// resolve via GetUserByToken. Both branches must reject
// with 401.
g.Use(func(c *gin.Context) {
ctx := c.Request.Context()
auth := c.GetHeader("Authorization")
if auth == "" {
common.ResponseWithCodeData(c, common.CodeUnauthorized, nil, "Authorization required")
c.Abort()
return
}
if u, code, err := stub.GetUserByToken(auth); err != nil || code != common.CodeSuccess {
if u, code, err := stub.GetUserByToken(ctx, auth); err != nil || code != common.CodeSuccess {
common.ResponseWithCodeData(c, common.CodeUnauthorized, nil, "Invalid auth credentials")
c.Abort()
return

8
internal/handler/chat_audio.go
View File

@@ -58,6 +58,7 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) {
return
}
ctx := c.Request.Context()
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, chatAudioSpeechMaxBodyBytes)
var req chatAudioSpeechRequest
@@ -84,7 +85,7 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) {
if seg == "" {
continue
}
resp, err := driver.AudioSpeech(&modelName, &seg, apiConfig, &modelModule.TTSConfig{Format: "mp3"}, nil)
resp, err := driver.AudioSpeech(ctx, &modelName, &seg, apiConfig, &modelModule.TTSConfig{Format: "mp3"}, nil)
if err != nil {
common.Warn("chat TTS synthesis failed",
zap.Int("segmentIndex", i),
@@ -154,6 +155,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) {
return
}
ctx := c.Request.Context()
// Cap the body before any multipart parsing so an oversized upload is
// rejected instead of being drained into memory or spooled to disk.
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, chatAudioUploadMaxBytes)
@@ -234,7 +236,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) {
return nil
}
if err = driver.TranscribeAudioWithSender(&modelName, &tmpPath, apiConfig, &modelModule.ASRConfig{}, nil, sender); err != nil {
if err = driver.TranscribeAudioWithSender(ctx, &modelName, &tmpPath, apiConfig, &modelModule.ASRConfig{}, nil, sender); err != nil {
errEvent := map[string]interface{}{"event": "error", "text": err.Error()}
data, _ := json.Marshal(errEvent)
_, _ = c.Writer.WriteString(fmt.Sprintf("data: %s\n\n", data))
@@ -252,7 +254,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) {
return
}
resp, err := driver.TranscribeAudio(&modelName, &tmpPath, apiConfig, &modelModule.ASRConfig{}, nil)
resp, err := driver.TranscribeAudio(ctx, &modelName, &tmpPath, apiConfig, &modelModule.ASRConfig{}, nil)
if err != nil {
common.ErrorWithCode(c, common.CodeServerError, err.Error())
return

3
internal/handler/chat_recommendation.go
View File

@@ -56,7 +56,8 @@ func (h *ChatHandler) Recommendation(c *gin.Context) {
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "question is required")
return
}
questions, err := service.GenerateRelatedQuestions(user.ID, req.Question, req.SearchID, h.searchSvc, h.tenantSvc, h.llm)
ctx := c.Request.Context()
questions, err := service.GenerateRelatedQuestions(ctx, user.ID, req.Question, req.SearchID, h.searchSvc, h.tenantSvc, h.llm)
if err != nil {
jsonInternalError(c, err)
return

6
internal/handler/chunk.go
View File

@@ -16,6 +16,7 @@
package handler
import (
"context"
"encoding/json"
"errors"
"fmt"
@@ -39,7 +40,7 @@ type chunkService interface {
UpdateChunk(req *service.UpdateChunkRequest, userID string) error
RemoveChunks(req *service.RemoveChunksRequest, userID string) (int64, error)
Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error)
AddChunk(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error)
}
@@ -698,6 +699,7 @@ func addChunkResponseMessage(code common.ErrorCode, err error) string {
}
func (h *ChunkHandler) AddChunk(c *gin.Context) {
ctx := c.Request.Context()
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.ErrorWithCode(c, errorCode, errorMessage)
@@ -756,7 +758,7 @@ func (h *ChunkHandler) AddChunk(c *gin.Context) {
ImageBase64: imageBase64,
}
resp, err := h.chunkService.AddChunk(&req, userID)
resp, err := h.chunkService.AddChunk(ctx, &req, userID)
if err != nil {
var codedErr service.ErrorCoder
if errors.As(err, &codedErr) {

17
internal/handler/chunk_test.go
View File

@@ -1,6 +1,7 @@
package handler
import (
"context"
"encoding/json"
"errors"
"net/http"
@@ -20,7 +21,7 @@ import (
// Only the methods actually called by the test are set; others panic.
type mockChunkSvc struct {
retrievalTestFn func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error)
addChunkFn func(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
addChunkFn func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error)
listFn func(req *service.ListChunksRequest, userID string) (*service.ListChunksResponse, error)
switchChunksFn func(userID, datasetID, documentID string, availableInt int, chunkIDs []string) error
updateChunkFn func(req *service.UpdateChunkRequest, userID string) error
@@ -82,9 +83,9 @@ func (m *mockChunkSvc) StopParsing(userID, datasetID string, req service.StopPar
func (m *mockChunkSvc) Parse(string, string, *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
panic("not implemented")
}
func (m *mockChunkSvc) AddChunk(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
func (m *mockChunkSvc) AddChunk(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
if m.addChunkFn != nil {
return m.addChunkFn(req, userID)
return m.addChunkFn(ctx, req, userID)
}
return &service.AddChunkResponse{Chunk: map[string]interface{}{"id": "chunk-1"}}, nil
}
@@ -641,7 +642,7 @@ func (e addChunkTestError) Code() common.ErrorCode { return e.code }
func TestChunkHandlerAddChunkSuccess(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
if userID != "user1" {
t.Fatalf("userID = %q, want user1", userID)
}
@@ -679,7 +680,7 @@ func TestChunkHandlerAddChunkSuccess(t *testing.T) {
func TestChunkHandlerAddChunkPathIDsOverrideBody(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
if req.DatasetID != "kb1" || req.DocumentID != "doc1" {
t.Fatalf("path IDs were not preserved: %#v", req)
}
@@ -710,7 +711,7 @@ func TestChunkHandlerAddChunkPathIDsOverrideBody(t *testing.T) {
func TestChunkHandlerAddChunkCodedError(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
return nil, addChunkTestError{code: common.CodeDataError, msg: "`content` is required"}
},
}
@@ -761,7 +762,7 @@ func TestChunkHandlerAddChunkValidatesListFields(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
t.Fatal("service should not be called for invalid request")
return nil, nil
},
@@ -792,7 +793,7 @@ func TestChunkHandlerAddChunkValidatesListFields(t *testing.T) {
func TestChunkHandlerAddChunkHidesServerErrorDetails(t *testing.T) {
mock := &mockChunkSvc{
addChunkFn: func(req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
addChunkFn: func(ctx context.Context, req *service.AddChunkRequest, userID string) (*service.AddChunkResponse, error) {
return nil, addChunkTestError{code: common.CodeServerError, msg: "encode chunk embedding: provider secret"}
},
}

3
internal/handler/dataset.go
View File

@@ -708,7 +708,8 @@ func (h *DatasetsHandler) CheckEmbedding(c *gin.Context) {
return
}
data, code, err := h.datasetsService.CheckEmbedding(userID, datasetID, &req)
ctx := c.Request.Context()
data, code, err := h.datasetsService.CheckEmbedding(ctx, userID, datasetID, &req)
if err != nil {
if code == common.CodeNotEffective {
common.ResponseWithCodeData(c, code, data, err.Error())

49
internal/handler/providers.go
View File

@@ -190,6 +190,7 @@ func (h *ProviderHandler) CreateProviderInstance(c *gin.Context) {
return
}
ctx := c.Request.Context()
var req CreateProviderInstanceRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ErrorWithCode(c, common.CodeBadRequest, err.Error())
@@ -211,7 +212,7 @@ func (h *ProviderHandler) CreateProviderInstance(c *gin.Context) {
return
}
_, err := h.modelProviderService.CreateProviderInstance(providerName, req.InstanceName, req.APIKey, req.BaseURL, req.Region, userID, req.ModelInfo)
_, err := h.modelProviderService.CreateProviderInstance(ctx, providerName, req.InstanceName, req.APIKey, req.BaseURL, req.Region, userID, req.ModelInfo)
if err != nil {
common.ErrorWithCode(c, common.CodeServerError, err.Error())
return
@@ -276,9 +277,10 @@ func (h *ProviderHandler) ShowInstanceBalance(c *gin.Context) {
}
userID := c.GetString("user_id")
ctx := c.Request.Context()
// Get tenant ID from user
balance, errorCode, err := h.modelProviderService.ShowInstanceBalance(providerName, instanceName, userID)
balance, errorCode, err := h.modelProviderService.ShowInstanceBalance(ctx, providerName, instanceName, userID)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -294,6 +296,7 @@ func (h *ProviderHandler) CheckConnection(c *gin.Context) {
return
}
ctx := c.Request.Context()
var req service.CheckConnectionRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, err.Error())
@@ -301,7 +304,7 @@ func (h *ProviderHandler) CheckConnection(c *gin.Context) {
}
userID := c.GetString("user_id")
errCode, err := h.modelProviderService.CheckConnection(providerName, req.APIKey, req.Region, req.BaseURL, req.InstanceID, userID, service.ListModelNames(req.ModelInfo))
errCode, err := h.modelProviderService.CheckConnection(ctx, providerName, req.APIKey, req.Region, req.BaseURL, req.InstanceID, userID, service.ListModelNames(req.ModelInfo))
if err != nil {
common.ErrorWithCode(c, errCode, err.Error())
return
@@ -311,6 +314,7 @@ func (h *ProviderHandler) CheckConnection(c *gin.Context) {
}
func (h *ProviderHandler) CheckInstanceConnection(c *gin.Context) {
ctx := c.Request.Context()
providerName := c.Param("provider_name")
if providerName == "" {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "Provider name is required")
@@ -336,7 +340,7 @@ func (h *ProviderHandler) CheckInstanceConnection(c *gin.Context) {
baseURL, _ := instanceInfo["base_url"].(string)
instanceID, _ := instanceInfo["id"].(string)
errorCode, err := h.modelProviderService.CheckConnection(providerName, apikey, region, baseURL, instanceID, userID, nil)
errorCode, err := h.modelProviderService.CheckConnection(ctx, providerName, apikey, region, baseURL, instanceID, userID, nil)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -346,6 +350,7 @@ func (h *ProviderHandler) CheckInstanceConnection(c *gin.Context) {
}
func (h *ProviderHandler) ListTasks(c *gin.Context) {
ctx := c.Request.Context()
providerName := c.Param("provider_name")
if providerName == "" {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "Provider name is required")
@@ -361,7 +366,7 @@ func (h *ProviderHandler) ListTasks(c *gin.Context) {
userID := c.GetString("user_id")
// Get tenant ID from user
listTaskResponse, errorCode, err := h.modelProviderService.ListTasks(providerName, instanceName, userID)
listTaskResponse, errorCode, err := h.modelProviderService.ListTasks(ctx, providerName, instanceName, userID)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -371,6 +376,7 @@ func (h *ProviderHandler) ListTasks(c *gin.Context) {
}
func (h *ProviderHandler) ShowTask(c *gin.Context) {
ctx := c.Request.Context()
providerName := c.Param("provider_name")
if providerName == "" {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "Provider name is required")
@@ -392,7 +398,7 @@ func (h *ProviderHandler) ShowTask(c *gin.Context) {
userID := c.GetString("user_id")
// Get tenant ID from user
taskResponse, errorCode, err := h.modelProviderService.ShowTask(providerName, instanceName, taskID, userID)
taskResponse, errorCode, err := h.modelProviderService.ShowTask(ctx, providerName, instanceName, taskID, userID)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -411,6 +417,7 @@ type AlterProviderInstanceRequest struct {
}
func (h *ProviderHandler) AlterProviderInstance(c *gin.Context) {
ctx := c.Request.Context()
providerName := c.Param("provider_name")
if providerName == "" {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "Provider name is required")
@@ -440,7 +447,7 @@ func (h *ProviderHandler) AlterProviderInstance(c *gin.Context) {
verify = *req.Verify
}
code, err := h.modelProviderService.AlterProviderInstance(userID, providerName, instanceName, req.InstanceName, req.APIKey, req.BaseURL, req.Region, req.ModelInfo, verify)
code, err := h.modelProviderService.AlterProviderInstance(ctx, userID, providerName, instanceName, req.InstanceName, req.APIKey, req.BaseURL, req.Region, req.ModelInfo, verify)
if err != nil {
common.ErrorWithCode(c, code, err.Error())
return
@@ -477,6 +484,7 @@ func (h *ProviderHandler) DropProviderInstance(c *gin.Context) {
}
func (h *ProviderHandler) ListInstanceModels(c *gin.Context) {
ctx := c.Request.Context()
providerName := c.Param("provider_name")
if providerName == "" {
common.ResponseWithHttpCodeData(c, http.StatusBadRequest, 400, nil, "Provider name is required")
@@ -498,7 +506,7 @@ func (h *ProviderHandler) ListInstanceModels(c *gin.Context) {
if keywords == "true" {
// list supported models
modelList, err := h.modelProviderService.ListSupportedModels(providerName, instanceName, c.GetString("user_id"))
modelList, err := h.modelProviderService.ListSupportedModels(ctx, providerName, instanceName, c.GetString("user_id"))
if err != nil {
common.ErrorWithCode(c, common.CodeServerError, err.Error())
return
@@ -694,6 +702,7 @@ type ChatToModelRequest struct {
}
func (h *ProviderHandler) ChatToModel(c *gin.Context) {
ctx := c.Request.Context()
var req ChatToModelRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -799,6 +808,7 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) {
// Stream response using sender function (the best performance, no channel)
errorCode, err := h.modelProviderService.ChatToModelStreamWithSender(
ctx,
req.ProviderName,
req.InstanceName,
req.ModelName,
@@ -830,6 +840,7 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) {
messages[i] = models.Message{Role: role, Content: content}
}
response, errorCode, err = h.modelProviderService.ChatToModelWithMessages(
ctx,
req.ProviderName,
req.InstanceName,
req.ModelName,
@@ -864,6 +875,7 @@ type EmbedTextRequest struct {
}
func (h *ProviderHandler) EmbedText(c *gin.Context) {
ctx := c.Request.Context()
var req EmbedTextRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -909,7 +921,7 @@ func (h *ProviderHandler) EmbedText(c *gin.Context) {
var errorCode common.ErrorCode
var err error
response, errorCode, err = h.modelProviderService.EmbedText(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Texts, &apiConfig, &embeddingConfig)
response, errorCode, err = h.modelProviderService.EmbedText(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Texts, &apiConfig, &embeddingConfig)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -929,6 +941,7 @@ type RerankDocumentRequest struct {
}
func (h *ProviderHandler) RerankDocument(c *gin.Context) {
ctx := c.Request.Context()
var req RerankDocumentRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -974,7 +987,7 @@ func (h *ProviderHandler) RerankDocument(c *gin.Context) {
var errorCode common.ErrorCode
var err error
response, errorCode, err = h.modelProviderService.RerankDocument(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Query, req.Documents, &apiConfig, &rerankConfig)
response, errorCode, err = h.modelProviderService.RerankDocument(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Query, req.Documents, &apiConfig, &rerankConfig)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -995,6 +1008,7 @@ type TranscribeAudioRequest struct {
}
func (h *ProviderHandler) TranscribeAudio(c *gin.Context) {
ctx := c.Request.Context()
var req TranscribeAudioRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -1070,7 +1084,7 @@ func (h *ProviderHandler) TranscribeAudio(c *gin.Context) {
}
// Stream response using sender function ( the best performance, no channel)
errorCode, err := h.modelProviderService.TranscribeAudioStream(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.File, &apiConfig, &asrConfig, sender)
errorCode, err := h.modelProviderService.TranscribeAudioStream(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.File, &apiConfig, &asrConfig, sender)
if errorCode != common.CodeSuccess {
c.SSEvent("error", err.Error())
}
@@ -1082,7 +1096,7 @@ func (h *ProviderHandler) TranscribeAudio(c *gin.Context) {
var errorCode common.ErrorCode
var err error
response, errorCode, err = h.modelProviderService.TranscribeAudio(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.File, &apiConfig, &asrConfig)
response, errorCode, err = h.modelProviderService.TranscribeAudio(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.File, &apiConfig, &asrConfig)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -1102,6 +1116,7 @@ type AudioSpeechRequest struct {
}
func (h *ProviderHandler) AudioSpeech(c *gin.Context) {
ctx := c.Request.Context()
var req AudioSpeechRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -1177,7 +1192,7 @@ func (h *ProviderHandler) AudioSpeech(c *gin.Context) {
}
// Stream response using sender function ( the best performance, no channel)
errorCode, err := h.modelProviderService.AudioSpeechStream(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Text, &apiConfig, &ttsConfig, sender)
errorCode, err := h.modelProviderService.AudioSpeechStream(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Text, &apiConfig, &ttsConfig, sender)
if errorCode != common.CodeSuccess {
c.SSEvent("error", err.Error())
}
@@ -1189,7 +1204,7 @@ func (h *ProviderHandler) AudioSpeech(c *gin.Context) {
var errorCode common.ErrorCode
var err error
response, errorCode, err = h.modelProviderService.AudioSpeech(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Text, &apiConfig, &ttsConfig)
response, errorCode, err = h.modelProviderService.AudioSpeech(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Text, &apiConfig, &ttsConfig)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -1208,6 +1223,7 @@ type OCRFileRequest struct {
}
func (h *ProviderHandler) OCRFile(c *gin.Context) {
ctx := c.Request.Context()
var req OCRFileRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -1251,7 +1267,7 @@ func (h *ProviderHandler) OCRFile(c *gin.Context) {
var errorCode common.ErrorCode
var err error
response, errorCode, err = h.modelProviderService.OCRFile(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Content, req.URL, &apiConfig, &OCRConfig)
response, errorCode, err = h.modelProviderService.OCRFile(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Content, req.URL, &apiConfig, &OCRConfig)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return
@@ -1270,6 +1286,7 @@ type ParseFileRequest struct {
}
func (h *ProviderHandler) ParseFile(c *gin.Context) {
ctx := c.Request.Context()
var req ParseFileRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -1313,7 +1330,7 @@ func (h *ProviderHandler) ParseFile(c *gin.Context) {
var errorCode common.ErrorCode
var err error
response, errorCode, err = h.modelProviderService.ParseFile(req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Content, req.URL, &apiConfig, &parseFileConfig)
response, errorCode, err = h.modelProviderService.ParseFile(ctx, req.ProviderName, req.InstanceName, req.ModelName, req.ModelID, userID, req.Content, req.URL, &apiConfig, &parseFileConfig)
if err != nil {
common.ErrorWithCode(c, errorCode, err.Error())
return

10
internal/handler/searchbot.go
View File

@@ -131,6 +131,7 @@ func (h *SearchBotHandler) SetAskService(svc *service.AskService) { h.askSvc = s
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/searchbots/related_questions [post]
func (h *SearchBotHandler) Handle(c *gin.Context) {
ctx := c.Request.Context()
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.ErrorWithCode(c, errorCode, errorMessage)
@@ -148,7 +149,7 @@ func (h *SearchBotHandler) Handle(c *gin.Context) {
return
}
questions, err := service.GenerateRelatedQuestions(user.ID, req.Question, req.SearchID, h.searchSvc, h.tenantSvc, h.llm)
questions, err := service.GenerateRelatedQuestions(ctx, user.ID, req.Question, req.SearchID, h.searchSvc, h.tenantSvc, h.llm)
if err != nil {
common.Warn("searchbot related questions failed", zap.String("error", err.Error()))
common.ResponseWithCodeData(c, common.CodeOperatingError, nil, "LLM call failed")
@@ -378,10 +379,11 @@ func (h *SearchBotHandler) MindMap(c *gin.Context) {
common.SuccessWithData(c, mindMap, "success")
}
// SearchbotDetail returns the public share-page bootstrap payload for a
// SearchBotDetail returns the public share-page bootstrap payload for a
// search app. The route is mounted under apiNoAuth but still requires a beta
// token, matching Python's AUTH_BETA flow.
func (h *SearchBotHandler) SearchbotDetail(c *gin.Context) {
func (h *SearchBotHandler) SearchBotDetail(c *gin.Context) {
ctx := c.Request.Context()
searchID := strings.TrimSpace(c.Query("search_id"))
if searchID == "" {
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "search_id is required")
@@ -389,7 +391,7 @@ func (h *SearchBotHandler) SearchbotDetail(c *gin.Context) {
}
userSvc := service.NewUserService()
user, code, err := userSvc.GetUserByBetaAPIToken(c.GetHeader("Authorization"))
user, code, err := userSvc.GetUserByBetaAPIToken(ctx, c.GetHeader("Authorization"))
if err != nil {
common.ResponseWithCodeData(c, code, nil, "Authentication error: API key is invalid!")
return

2
internal/handler/searchbot_detail_test.go
View File

@@ -58,7 +58,7 @@ func newSearchbotDetailRouter() *gin.Engine {
gin.SetMode(gin.TestMode)
h := NewSearchBotHandler(service.NewSearchService(), nil, nil, nil)
r := gin.New()
r.GET("/api/v1/searchbots/detail", h.SearchbotDetail)
r.GET("/api/v1/searchbots/detail", h.SearchBotDetail)
return r
}

32
internal/handler/user.go
View File

@@ -56,6 +56,7 @@ func NewUserHandler(userService *service.UserService) *UserHandler {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/users [post]
func (h *UserHandler) Register(c *gin.Context) {
ctx := c.Request.Context()
var req service.RegisterRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ResponseWithCodeData(c, common.CodeBadRequest, false, err.Error())
@@ -90,7 +91,7 @@ func (h *UserHandler) Register(c *gin.Context) {
c.Header("Access-Control-Allow-Headers", "*")
c.Header("Access-Control-Expose-Headers", "Authorization")
profile := h.userService.GetUserProfile(user)
profile := h.userService.GetUserProfile(ctx, user)
common.SuccessWithData(c, profile, fmt.Sprintf("%s, welcome aboard!", req.Nickname))
}
@@ -104,6 +105,7 @@ func (h *UserHandler) Register(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/users/login [post]
func (h *UserHandler) Login(c *gin.Context) {
ctx := c.Request.Context()
startAt := time.Now()
operationLog := &common.OperationLog{
EventTime: startAt,
@@ -170,7 +172,7 @@ func (h *UserHandler) Login(c *gin.Context) {
c.Header("Access-Control-Allow-Headers", "*")
c.Header("Access-Control-Expose-Headers", "Authorization")
profile := h.userService.GetUserProfile(user)
profile := h.userService.GetUserProfile(ctx, user)
common.SuccessWithData(c, profile, "Welcome back!")
}
@@ -184,6 +186,7 @@ func (h *UserHandler) Login(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /v1/user/login [post]
func (h *UserHandler) LoginByEmail(c *gin.Context) {
ctx := c.Request.Context()
startAt := time.Now()
operationLog := &common.OperationLog{
EventTime: startAt,
@@ -221,7 +224,7 @@ func (h *UserHandler) LoginByEmail(c *gin.Context) {
return
}
user, code, err := h.userService.LoginByEmail(&req)
user, code, err := h.userService.LoginByEmail(ctx, &req)
if err != nil {
common.ResponseWithCodeData(c, code, false, err.Error())
operationLog.ErrorCode = uint16(code)
@@ -255,7 +258,7 @@ func (h *UserHandler) LoginByEmail(c *gin.Context) {
c.Header("Access-Control-Allow-Headers", "*")
c.Header("Access-Control-Expose-Headers", "Authorization")
profile := h.userService.GetUserProfile(user)
profile := h.userService.GetUserProfile(ctx, user)
common.SuccessWithData(c, profile, "Welcome back!")
}
@@ -269,6 +272,7 @@ func (h *UserHandler) LoginByEmail(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/users/{id} [get]
func (h *UserHandler) GetUserByID(c *gin.Context) {
ctx := c.Request.Context()
idStr := c.Param("id")
id, err := strconv.ParseUint(idStr, 10, 32)
if err != nil {
@@ -276,7 +280,7 @@ func (h *UserHandler) GetUserByID(c *gin.Context) {
return
}
user, code, err := h.userService.GetUserByID(uint(id))
user, code, err := h.userService.GetUserByID(ctx, uint(id))
if err != nil {
common.ResponseWithCodeData(c, code, false, err.Error())
return
@@ -295,6 +299,7 @@ func (h *UserHandler) GetUserByID(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /v1/user/logout [post]
func (h *UserHandler) Logout(c *gin.Context) {
ctx := c.Request.Context()
startAt := time.Now()
operationLog := &common.OperationLog{
EventTime: startAt,
@@ -335,7 +340,7 @@ func (h *UserHandler) Logout(c *gin.Context) {
}
// Get user by access token
user, code, err := h.userService.GetUserByToken(token)
user, code, err := h.userService.GetUserByToken(ctx, token)
if err != nil {
common.ResponseWithHttpCodeData(c, http.StatusUnauthorized, code, nil, "Invalid access token")
c.Abort()
@@ -369,6 +374,7 @@ func (h *UserHandler) Logout(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /v1/user/info [get]
func (h *UserHandler) Info(c *gin.Context) {
ctx := c.Request.Context()
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.ErrorWithCode(c, errorCode, errorMessage)
@@ -376,7 +382,7 @@ func (h *UserHandler) Info(c *gin.Context) {
}
// Get user profile
profile := h.userService.GetUserProfile(user)
profile := h.userService.GetUserProfile(ctx, user)
common.SuccessWithData(c, profile, "success")
}
@@ -392,6 +398,7 @@ func (h *UserHandler) Info(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/users/me [patch]
func (h *UserHandler) Setting(c *gin.Context) {
ctx := c.Request.Context()
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.ErrorWithCode(c, errorCode, errorMessage)
@@ -406,7 +413,7 @@ func (h *UserHandler) Setting(c *gin.Context) {
}
// Update user settings
code, err := h.userService.UpdateUserSettings(user, &req)
code, err := h.userService.UpdateUserSettings(ctx, user, &req)
if err != nil {
if code == common.CodeExceptionError {
common.ResponseWithCodeData(c, common.CodeExceptionError, false, err.Error())
@@ -430,6 +437,7 @@ func (h *UserHandler) Setting(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /v1/user/setting/password [post]
func (h *UserHandler) ChangePassword(c *gin.Context) {
ctx := c.Request.Context()
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.ErrorWithCode(c, errorCode, errorMessage)
@@ -444,7 +452,7 @@ func (h *UserHandler) ChangePassword(c *gin.Context) {
}
// Change password
code, err := h.userService.ChangePassword(user, &req)
code, err := h.userService.ChangePassword(ctx, user, &req)
if err != nil {
common.ResponseWithCodeData(c, code, false, err.Error())
return
@@ -482,6 +490,7 @@ func (h *UserHandler) GetLoginChannels(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /v1/user/set_tenant_info [post]
func (h *UserHandler) SetTenantInfo(c *gin.Context) {
ctx := c.Request.Context()
user, errorCode, errorMessage := GetUser(c)
if errorCode != common.CodeSuccess {
common.ErrorWithCode(c, errorCode, errorMessage)
@@ -531,7 +540,7 @@ func (h *UserHandler) SetTenantInfo(c *gin.Context) {
req.TTSID = &value
}
code, err := h.userService.SetTenantInfo(user.ID, &req)
code, err := h.userService.SetTenantInfo(ctx, user.ID, &req)
if err != nil {
common.ResponseWithCodeData(c, code, nil, err.Error())
return
@@ -684,6 +693,7 @@ func (h *UserHandler) ForgotVerifyOTP(c *gin.Context) {
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/auth/password/reset [post]
func (h *UserHandler) ForgotResetPassword(c *gin.Context) {
ctx := c.Request.Context()
var req service.ForgotResetPasswordRequest
if err := c.ShouldBindJSON(&req); err != nil {
common.ResponseWithCodeData(c, common.CodeArgumentError, false, err.Error())
@@ -714,7 +724,7 @@ func (h *UserHandler) ForgotResetPassword(c *gin.Context) {
// already in the Authorization header). Mirror the Python contract
// `user.to_safe_dict(for_self=True)` by stripping those fields before
// writing. PR #15290 review.
profile := h.userService.GetUserProfile(user)
profile := h.userService.GetUserProfile(ctx, user)
delete(profile, "password")
delete(profile, "access_token")
common.SuccessWithData(c, profile, "Password reset successful. Logged in.")