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

Go: add context, part13 (#17445)

Browse Source 
Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-07-27 21:19:39 +08:00
committed by GitHub
parent 1cce9b64f1
commit 1436fcaca5
31 changed files with 536 additions and 489 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/agent/component/agent.go
View File

@@ -753,7 +753,7 @@ func (c *AgentComponent) invokeNow(ctx context.Context, db *gorm.DB, inputs map[
}
var err error
p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, p.ModelID, p.Driver, p.APIKey, p.BaseURL)
p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, db, p.ModelID, p.Driver, p.APIKey, p.BaseURL)
if err != nil {
return nil, err
}

48
internal/agent/component/browser.go
View File

@@ -240,7 +240,7 @@ func (p *browserParam) AsDict() map[string]any {
}
// BrowserComponent is the canvas Browser node. Owns its static
// param; delegates the multi-step agent run to StagehandInvoker.
// param; delegates the multistep agent run to StagehandInvoker.
type BrowserComponent struct {
name string
param browserParam
@@ -250,10 +250,10 @@ type BrowserComponent struct {
func NewBrowserComponent(params map[string]any) (Component, error) {
p := &browserParam{}
if err := p.Update(params); err != nil {
return nil, fmt.Errorf("Browser: param update: %w", err)
return nil, fmt.Errorf("browser: param update: %w", err)
}
if err := p.Check(); err != nil {
return nil, fmt.Errorf("Browser: param check: %w", err)
return nil, fmt.Errorf("browser: param check: %w", err)
}
return &BrowserComponent{
name: componentNameBrowser,
@@ -285,15 +285,15 @@ func (b *BrowserComponent) Name() string { return b.name }
func (b *BrowserComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) {
state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx)
if err != nil {
return nil, fmt.Errorf("Browser: %w", err)
return nil, fmt.Errorf("browser: %w", err)
}
if state == nil {
return nil, errors.New("Browser: nil canvas state")
return nil, errors.New("browser: nil canvas state")
}
tenantID, _ := state.Sys["tenant_id"].(string)
if tenantID == "" {
return nil, errors.New("Browser: tenant_id missing from canvas state (state.Sys[\"tenant_id\"])")
return nil, errors.New("browser: tenant_id missing from canvas state (state.Sys[\"tenant_id\"])")
}
// 1. Resolve prompts template.
@@ -303,13 +303,13 @@ func (b *BrowserComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[s
}
resolvedPrompts, err := runtime.ResolveTemplate(prompts, state)
if err != nil {
return nil, fmt.Errorf("Browser: resolve prompts template: %w", err)
return nil, fmt.Errorf("browser: resolve prompts template: %w", err)
}
// 2. Look up tenant model config.
providerName, modelName, apiKey, baseURL, err := resolveBrowserLLM(tenantID, b.param.LLMID)
providerName, modelName, apiKey, baseURL, err := resolveBrowserLLM(ctx, db, tenantID, b.param.LLMID)
if err != nil {
return nil, fmt.Errorf("Browser: tenant llm lookup (%q): %w", b.param.LLMID, err)
return nil, fmt.Errorf("browser: tenant llm lookup (%q): %w", b.param.LLMID, err)
}
baseURL = strings.TrimSpace(baseURL)
@@ -329,14 +329,14 @@ func (b *BrowserComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[s
invoker := getDefaultStagehandInvoker()
rawJSON, err := invoker.RunExtract(ctx, req)
if err != nil {
return nil, fmt.Errorf("Browser: stagehand extract (model=%q, base_url=%s): %w",
return nil, fmt.Errorf("browser: stagehand extract (model=%q, base_url=%s): %w",
req.ModelName, browserBaseURLForLog(req.BaseURL), err)
}
// 5. Unmarshal the JSON-string result to get the plain text.
var content string
if err := json.Unmarshal([]byte(rawJSON), &content); err != nil {
return nil, fmt.Errorf("Browser: unmarshal extract result: %w", err)
if err = json.Unmarshal([]byte(rawJSON), &content); err != nil {
return nil, fmt.Errorf("browser: unmarshal extract result: %w", err)
}
// 6. Build the output map.
@@ -444,7 +444,7 @@ func (b *BrowserComponent) Outputs() map[string]string {
// Tests override the lookup via `tenantLLMLookupForTest` (a
// package-level function variable) so they don't need a real DB.
// Production code leaves the variable unset.
func resolveBrowserLLM(tenantID, llmID string) (providerName, modelName, apiKey, baseURL string, err error) {
func resolveBrowserLLM(ctx context.Context, db *gorm.DB, tenantID, llmID string) (providerName, modelName, apiKey, baseURL string, err error) {
if tenantLLMLookupForTest != nil {
oldModelName, factory := resolveLLMID(llmID)
apiKey, baseURL, err = tenantLLMLookupForTest(tenantID, oldModelName, factory)
@@ -452,7 +452,7 @@ func resolveBrowserLLM(tenantID, llmID string) (providerName, modelName, apiKey,
return factory, oldModelName, apiKey, baseURL, err
}
providerName, modelName, apiKey, baseURL, err = resolveTenantModelBrowserLLM(tenantID, llmID)
providerName, modelName, apiKey, baseURL, err = resolveTenantModelBrowserLLM(ctx, db, tenantID, llmID)
if err == nil {
baseURL = browserOpenAICompatibleBaseURL(baseURL, providerName)
return providerName, modelName, apiKey, baseURL, nil
@@ -460,7 +460,7 @@ func resolveBrowserLLM(tenantID, llmID string) (providerName, modelName, apiKey,
modelErr := err
oldModelName, factory := resolveLLMID(llmID)
apiKey, baseURL, oldErr := resolveTenantLLM(tenantID, oldModelName, factory)
apiKey, baseURL, oldErr := resolveTenantLLM(ctx, db, tenantID, oldModelName, factory)
if oldErr == nil {
baseURL = browserOpenAICompatibleBaseURL(baseURL, factory)
return factory, oldModelName, apiKey, baseURL, nil
@@ -468,8 +468,8 @@ func resolveBrowserLLM(tenantID, llmID string) (providerName, modelName, apiKey,
return "", "", "", "", fmt.Errorf("tenant_model lookup: %v; tenant_llm fallback: %w", modelErr, oldErr)
}
func resolveTenantModelBrowserLLM(tenantID, modelID string) (providerName, modelName, apiKey, baseURL string, err error) {
modelRow, err := dao.NewTenantModelDAO().GetByID(modelID)
func resolveTenantModelBrowserLLM(ctx context.Context, db *gorm.DB, tenantID, modelID string) (providerName, modelName, apiKey, baseURL string, err error) {
modelRow, err := dao.NewTenantModelDAO().GetByID(ctx, db, modelID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return "", "", "", "", err
@@ -483,7 +483,7 @@ func resolveTenantModelBrowserLLM(tenantID, modelID string) (providerName, model
return "", "", "", "", fmt.Errorf("tenant model id=%s cannot be used as %s model", modelID, entity.ModelTypeChat.String())
}
provider, err := dao.NewTenantModelProviderDAO().GetByID(modelRow.ProviderID)
provider, err := dao.NewTenantModelProviderDAO().GetByID(ctx, db, modelRow.ProviderID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return "", "", "", "", fmt.Errorf("provider id=%s not found for model id=%s", modelRow.ProviderID, modelID)
@@ -497,7 +497,7 @@ func resolveTenantModelBrowserLLM(tenantID, modelID string) (providerName, model
return "", "", "", "", fmt.Errorf("tenant %s has no access to provider owned by tenant %s", tenantID, provider.TenantID)
}
instance, err := dao.NewTenantModelInstanceDAO().GetByID(modelRow.InstanceID)
instance, err := dao.NewTenantModelInstanceDAO().GetByID(ctx, db, modelRow.InstanceID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return "", "", "", "", fmt.Errorf("instance id=%s not found for model id=%s", modelRow.InstanceID, modelID)
@@ -511,7 +511,7 @@ func resolveTenantModelBrowserLLM(tenantID, modelID string) (providerName, model
apiKey = instance.APIKey
if strings.TrimSpace(instance.Extra) != "" {
var extra map[string]string
if err := json.Unmarshal([]byte(instance.Extra), &extra); err != nil {
if err = json.Unmarshal([]byte(instance.Extra), &extra); err != nil {
return "", "", "", "", err
}
baseURL = extra["base_url"]
@@ -537,18 +537,18 @@ func browserOpenAICompatibleBaseURL(baseURL, provider string) string {
//
// TODO(v2): this helper can move to `internal/dao` so the LLM
// component (`llm.go`) and other future components can share it.
func resolveTenantLLM(tenantID, modelName, factory string) (apiKey, baseURL string, err error) {
dao := dao.NewTenantLLMDAO()
func resolveTenantLLM(ctx context.Context, db *gorm.DB, tenantID, modelName, factory string) (apiKey, baseURL string, err error) {
llmDAO := dao.NewTenantLLMDAO()
var (
row *entity.TenantLLM
)
if factory != "" {
row, err = dao.GetByTenantFactoryAndModelName(tenantID, factory, modelName)
row, err = llmDAO.GetByTenantFactoryAndModelName(ctx, db, tenantID, factory, modelName)
} else {
// No factory suffix on llm_id; fall back to a single-key
// lookup (errors if the model is registered under multiple
// factories — caller must use the explicit form).
row, err = dao.GetByTenantAndModelName(tenantID, "", modelName)
row, err = llmDAO.GetByTenantAndModelName(ctx, db, tenantID, "", modelName)
}
if err != nil {
return "", "", err

3
internal/agent/component/browser_test.go
View File

@@ -249,7 +249,8 @@ func TestResolveBrowserLLM_ResolvesTenantModelID(t *testing.T) {
tenantLLMLookupForTest = nil
t.Cleanup(func() { tenantLLMLookupForTest = prevLookup })
provider, model, apiKey, baseURL, err := resolveBrowserLLM("tenant-1", "tenant-model-1")
ctx := t.Context()
provider, model, apiKey, baseURL, err := resolveBrowserLLM(ctx, db, "tenant-1", "tenant-model-1")
if err != nil {
t.Fatalf("resolveBrowserLLM: %v", err)
}

2
internal/agent/component/categorize.go
View File

@@ -68,7 +68,7 @@ func (c *CategorizeComponent) Name() string { return "Categorize" }
func (c *CategorizeComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) {
p := mergeCategorizeParam(c.param, inputs)
var err error
p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, p.ModelID, p.Driver, p.APIKey, p.BaseURL)
p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, db, p.ModelID, p.Driver, p.APIKey, p.BaseURL)
if err != nil {
return nil, err
}

2
internal/agent/component/llm.go
View File

@@ -321,7 +321,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[strin
// reference selected in the agent canvas is passed verbatim to the LLM
// driver, causing 400s for custom-added models.
var err error
p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, p.ModelID, p.Driver, p.APIKey, p.BaseURL)
p.ModelID, p.Driver, p.APIKey, p.BaseURL, err = resolveChatModelRef(ctx, db, p.ModelID, p.Driver, p.APIKey, p.BaseURL)
if err != nil {
return nil, fmt.Errorf("component: LLM.Invoke: resolve model: %w", err)
}

50
internal/agent/component/llm_credentials.go
View File

@@ -20,7 +20,7 @@ import (
// driver/model pair when the canvas DSL omitted them. It first checks the old
// tenant_llm table, then falls back to tenant_model_provider +
// tenant_model_instance when the composite llm_id carries an instance name.
func resolveTenantLLMConfig(ctx context.Context, driver, modelID, apiKey, baseURL, originalModelID string) (string, string) {
func resolveTenantLLMConfig(ctx context.Context, db *gorm.DB, driver, modelID, apiKey, baseURL, originalModelID string) (string, string) {
if apiKey != "" || driver == "" || modelID == "" {
return apiKey, baseURL
}
@@ -35,19 +35,19 @@ func resolveTenantLLMConfig(ctx context.Context, driver, modelID, apiKey, baseUR
return apiKey, baseURL
}
if resolvedKey, resolvedBaseURL, ok := resolveTenantLLMCredentials(tid, driver, modelID, baseURL); ok {
if resolvedKey, resolvedBaseURL, ok := resolveTenantLLMCredentials(ctx, db, tid, driver, modelID, baseURL); ok {
return resolvedKey, resolvedBaseURL
}
if originalModelID == "" {
return apiKey, baseURL
}
if resolvedKey, resolvedBaseURL, ok := resolveTenantModelInstanceCredentials(tid, originalModelID, baseURL); ok {
if resolvedKey, resolvedBaseURL, ok := resolveTenantModelInstanceCredentials(ctx, db, tid, originalModelID, baseURL); ok {
return resolvedKey, resolvedBaseURL
}
return apiKey, baseURL
}
func resolveChatModelRef(ctx context.Context, modelID, driver, apiKey, baseURL string) (string, string, string, string, error) {
func resolveChatModelRef(ctx context.Context, db *gorm.DB, modelID, driver, apiKey, baseURL string) (string, string, string, string, error) {
originalModelID := modelID
if driver == "" && modelID != "" {
if m, prov, ok := agentProviderLastSegmentSplit(modelID); ok {
@@ -56,7 +56,7 @@ func resolveChatModelRef(ctx context.Context, modelID, driver, apiKey, baseURL s
}
}
if driver == "" && modelID != "" {
resolvedModelID, resolvedDriver, resolvedAPIKey, resolvedBaseURL, ok, err := resolveTenantChatModelByID(ctx, modelID, apiKey, baseURL)
resolvedModelID, resolvedDriver, resolvedAPIKey, resolvedBaseURL, ok, err := resolveTenantChatModelByID(ctx, db, modelID, apiKey, baseURL)
if err != nil {
return "", "", "", "", err
}
@@ -67,15 +67,15 @@ func resolveChatModelRef(ctx context.Context, modelID, driver, apiKey, baseURL s
baseURL = resolvedBaseURL
}
}
apiKey, baseURL = resolveTenantLLMConfig(ctx, driver, modelID, apiKey, baseURL, originalModelID)
apiKey, baseURL = resolveTenantLLMConfig(ctx, db, driver, modelID, apiKey, baseURL, originalModelID)
return modelID, driver, apiKey, baseURL, nil
}
// resolveTenantLLMCredentials looks up the old tenant_llm table for the given
// tenant / factory / model. Returns true when credentials were found.
func resolveTenantLLMCredentials(tid, driver, modelID, baseURL string) (string, string, bool) {
func resolveTenantLLMCredentials(ctx context.Context, db *gorm.DB, tid, driver, modelID, baseURL string) (string, string, bool) {
common.Debug("llm credentials: tenant_llm lookup", zap.String("tid", tid), zap.String("factory", driver), zap.String("model", modelID))
row, err := dao.NewTenantLLMDAO().GetByTenantFactoryAndModelName(tid, driver, modelID)
row, err := dao.NewTenantLLMDAO().GetByTenantFactoryAndModelName(ctx, db, tid, driver, modelID)
if err != nil {
common.Debug("llm credentials: tenant_llm lookup", zap.Error(err))
return "", baseURL, false
@@ -101,7 +101,7 @@ func resolveTenantLLMCredentials(tid, driver, modelID, baseURL string) (string,
// resolveTenantModelInstanceCredentials attempts to resolve llm credentials
// through tenant_model_provider + tenant_model_instance using the original
// composite llm_id (which still carries the instance name).
func resolveTenantModelInstanceCredentials(tid, compositeLLMID, baseURL string) (string, string, bool) {
func resolveTenantModelInstanceCredentials(ctx context.Context, db *gorm.DB, tid, compositeLLMID, baseURL string) (string, string, bool) {
modelName, instanceName, providerName := parseLLMIDParts(compositeLLMID)
if instanceName == "" {
common.Debug("llm credentials: new-table fallback skipped: no instance name", zap.String("composite_llm_id", compositeLLMID))
@@ -114,16 +114,16 @@ func resolveTenantModelInstanceCredentials(tid, compositeLLMID, baseURL string)
zap.String("model", modelName),
zap.String("instance", instanceName))
provider, err := dao.NewTenantModelProviderDAO().GetByTenantIDAndProviderName(tid, providerName)
provider, err := dao.NewTenantModelProviderDAO().GetByTenantIDAndProviderName(ctx, db, tid, providerName)
if err != nil || provider == nil {
common.Debug("llm credentials: new-table fallback: provider not found", zap.String("provider", providerName), zap.Error(err))
return "", baseURL, false
}
instance, err := dao.NewTenantModelInstanceDAO().GetByProviderIDAndInstanceName(provider.ID, instanceName)
instance, err := dao.NewTenantModelInstanceDAO().GetByProviderIDAndInstanceName(ctx, db, provider.ID, instanceName)
if err != nil || instance == nil {
if instanceName == "default" {
if fallback := findSoleActiveProviderInstance(provider.ID); fallback != nil {
if fallback := findSoleActiveProviderInstance(ctx, db, provider.ID); fallback != nil {
common.Debug("llm credentials: new-table fallback: remapped default instance to sole active instance",
zap.String("instance", fallback.InstanceName),
zap.String("provider", providerName))
@@ -158,8 +158,8 @@ func resolveTenantModelInstanceCredentials(tid, compositeLLMID, baseURL string)
return apiKey, baseURL, apiKey != ""
}
func findSoleActiveProviderInstance(providerID string) *entity.TenantModelInstance {
instances, err := dao.NewTenantModelInstanceDAO().GetAllInstancesByProviderID(providerID)
func findSoleActiveProviderInstance(ctx context.Context, db *gorm.DB, providerID string) *entity.TenantModelInstance {
instances, err := dao.NewTenantModelInstanceDAO().GetAllInstancesByProviderID(ctx, db, providerID)
if err != nil {
common.Debug("llm credentials: list provider instances", zap.Error(err))
return nil
@@ -180,7 +180,7 @@ func findSoleActiveProviderInstance(providerID string) *entity.TenantModelInstan
return active[0]
}
func resolveTenantChatModelByID(ctx context.Context, modelRef, apiKey, baseURL string) (string, string, string, string, bool, error) {
func resolveTenantChatModelByID(ctx context.Context, db *gorm.DB, modelRef, apiKey, baseURL string) (string, string, string, string, bool, error) {
if !isBareTenantModelID(modelRef) {
return "", "", apiKey, baseURL, false, nil
}
@@ -193,7 +193,7 @@ func resolveTenantChatModelByID(ctx context.Context, modelRef, apiKey, baseURL s
return "", "", apiKey, baseURL, false, nil
}
modelName, provider, modelKey, modelBaseURL, ok, err := resolveTenantChatModelByTenantModelID(tid, modelRef, apiKey, baseURL)
modelName, provider, modelKey, modelBaseURL, ok, err := resolveTenantChatModelByTenantModelID(ctx, db, tid, modelRef, apiKey, baseURL)
if err != nil {
return "", "", apiKey, baseURL, false, err
}
@@ -201,7 +201,7 @@ func resolveTenantChatModelByID(ctx context.Context, modelRef, apiKey, baseURL s
return modelName, provider, modelKey, modelBaseURL, true, nil
}
modelName, provider, instanceKey, instanceBaseURL, ok, err := resolveTenantChatModelByInstanceID(tid, modelRef, apiKey, baseURL)
modelName, provider, instanceKey, instanceBaseURL, ok, err := resolveTenantChatModelByInstanceID(ctx, db, tid, modelRef, apiKey, baseURL)
if err != nil {
return "", "", apiKey, baseURL, false, err
}
@@ -211,8 +211,8 @@ func resolveTenantChatModelByID(ctx context.Context, modelRef, apiKey, baseURL s
return "", "", apiKey, baseURL, false, fmt.Errorf("tenant chat model id %q not found", modelRef)
}
func resolveTenantChatModelByTenantModelID(tid, modelID, apiKey, baseURL string) (string, string, string, string, bool, error) {
model, err := dao.NewTenantModelDAO().GetByID(modelID)
func resolveTenantChatModelByTenantModelID(ctx context.Context, db *gorm.DB, tid, modelID, apiKey, baseURL string) (string, string, string, string, bool, error) {
model, err := dao.NewTenantModelDAO().GetByID(ctx, db, modelID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return "", "", apiKey, baseURL, false, nil
@@ -226,7 +226,7 @@ func resolveTenantChatModelByTenantModelID(tid, modelID, apiKey, baseURL string)
return "", "", apiKey, baseURL, false, fmt.Errorf("tenant model id %s cannot be used as chat model", modelID)
}
provider, err := dao.NewTenantModelProviderDAO().GetByID(model.ProviderID)
provider, err := dao.NewTenantModelProviderDAO().GetByID(ctx, db, model.ProviderID)
if err != nil {
return "", "", apiKey, baseURL, false, fmt.Errorf("resolve provider for tenant model id %q: %w", modelID, err)
}
@@ -234,7 +234,7 @@ func resolveTenantChatModelByTenantModelID(tid, modelID, apiKey, baseURL string)
return "", "", apiKey, baseURL, false, fmt.Errorf("tenant %s has no access to model id %s", tid, modelID)
}
instance, err := dao.NewTenantModelInstanceDAO().GetByID(model.InstanceID)
instance, err := dao.NewTenantModelInstanceDAO().GetByID(ctx, db, model.InstanceID)
if err != nil {
return "", "", apiKey, baseURL, false, fmt.Errorf("resolve instance for tenant model id %q: %w", modelID, err)
}
@@ -242,15 +242,15 @@ func resolveTenantChatModelByTenantModelID(tid, modelID, apiKey, baseURL string)
return model.ModelName, provider.ProviderName, apiKey, baseURL, true, nil
}
func resolveTenantChatModelByInstanceID(tid, instanceID, apiKey, baseURL string) (string, string, string, string, bool, error) {
instance, err := dao.NewTenantModelInstanceDAO().GetByID(instanceID)
func resolveTenantChatModelByInstanceID(ctx context.Context, db *gorm.DB, tid, instanceID, apiKey, baseURL string) (string, string, string, string, bool, error) {
instance, err := dao.NewTenantModelInstanceDAO().GetByID(ctx, db, instanceID)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return "", "", apiKey, baseURL, false, nil
}
return "", "", apiKey, baseURL, false, fmt.Errorf("resolve tenant model instance id %q: %w", instanceID, err)
}
provider, err := dao.NewTenantModelProviderDAO().GetByID(instance.ProviderID)
provider, err := dao.NewTenantModelProviderDAO().GetByID(ctx, db, instance.ProviderID)
if err != nil {
return "", "", apiKey, baseURL, false, fmt.Errorf("resolve provider for tenant model instance id %q: %w", instanceID, err)
}
@@ -258,7 +258,7 @@ func resolveTenantChatModelByInstanceID(tid, instanceID, apiKey, baseURL string)
return "", "", apiKey, baseURL, false, fmt.Errorf("tenant %s has no access to model instance id %s", tid, instanceID)
}
models, err := dao.NewTenantModelDAO().GetModelsByInstanceID(instance.ID)
models, err := dao.NewTenantModelDAO().GetModelsByInstanceID(ctx, db, instance.ID)
if err != nil {
return "", "", apiKey, baseURL, false, fmt.Errorf("resolve models for tenant model instance id %q: %w", instanceID, err)
}