diff --git a/internal/agent/component/agent.go b/internal/agent/component/agent.go index d4b117c2ca..db565e538b 100644 --- a/internal/agent/component/agent.go +++ b/internal/agent/component/agent.go @@ -464,6 +464,18 @@ func (c *AgentComponent) Invoke(ctx context.Context, inputs map[string]any) (map p.ModelID = m } } + if p.Driver == "" && p.ModelID != "" { + modelID, driver, apiKey, baseURL, ok, err := resolveTenantChatModelByID(ctx, p.ModelID, p.APIKey, p.BaseURL) + if err != nil { + return nil, err + } + if ok { + p.ModelID = modelID + p.Driver = driver + p.APIKey = apiKey + p.BaseURL = baseURL + } + } p.APIKey, p.BaseURL = resolveTenantLLMConfig(ctx, p.Driver, p.ModelID, p.APIKey, p.BaseURL, originalModelID) var state *runtime.CanvasState diff --git a/internal/agent/component/data_operations.go b/internal/agent/component/data_operations.go index 235a110d72..d3ccec3dcd 100644 --- a/internal/agent/component/data_operations.go +++ b/internal/agent/component/data_operations.go @@ -269,6 +269,50 @@ func (d *DataOperationsComponent) Inputs() map[string]string { } } +func (d *DataOperationsComponent) GetInputForm() map[string]any { + res := make(map[string]any) + for _, input := range d.param.Query { + for k, v := range getInputElementsFromText(input) { + res[k] = map[string]any{ + "name": v["name"], + "type": "line", + } + } + } + return map[string]any{} +} + +// TODO duplicate with string_transform.go, make it to a common method later --Haruko386 +// I can't do this since that PR not merged +func getInputElementsFromText(txt string) map[string]map[string]any { + res := make(map[string]map[string]any) + for _, match := range runtime.VarRefPattern.FindAllStringSubmatch(txt, -1) { + if len(match) < 2 { + continue + } + exp := match[1] + if _, ok := res[exp]; ok { + continue + } + cpnID, varName := "", exp + if at := strings.Index(exp, "@"); at > 0 { + cpnID = exp[:at] + varName = exp[at+1:] + } + name := exp + if cpnID != "" { + name = cpnID + "@" + varName + } + res[exp] = map[string]any{ + "name": name, + "value": nil, + "_retrieval": nil, + "_cpn_id": cpnID, + } + } + return res +} + // Outputs returns the transformed payload. func (d *DataOperationsComponent) Outputs() map[string]string { return map[string]string{ diff --git a/internal/agent/component/llm_credentials.go b/internal/agent/component/llm_credentials.go index 03ed1fd0ba..f1391af3bd 100644 --- a/internal/agent/component/llm_credentials.go +++ b/internal/agent/component/llm_credentials.go @@ -3,6 +3,8 @@ package component import ( "context" "encoding/json" + "errors" + "fmt" "strings" "ragflow/internal/agent/runtime" @@ -11,6 +13,7 @@ import ( "ragflow/internal/entity" "go.uber.org/zap" + "gorm.io/gorm" ) // resolveTenantLLMConfig fills tenant-scoped API credentials for the supplied @@ -153,6 +156,144 @@ func findSoleActiveProviderInstance(providerID string) *entity.TenantModelInstan return active[0] } +func resolveTenantChatModelByID(ctx context.Context, modelRef, apiKey, baseURL string) (string, string, string, string, bool, error) { + if !isBareTenantModelID(modelRef) { + return "", "", apiKey, baseURL, false, nil + } + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return "", "", apiKey, baseURL, false, nil + } + tid, _ := state.Sys["tenant_id"].(string) + if tid == "" { + return "", "", apiKey, baseURL, false, nil + } + + modelName, provider, modelKey, modelBaseURL, ok, err := resolveTenantChatModelByTenantModelID(tid, modelRef, apiKey, baseURL) + if err != nil { + return "", "", apiKey, baseURL, false, err + } + if ok { + return modelName, provider, modelKey, modelBaseURL, true, nil + } + + modelName, provider, instanceKey, instanceBaseURL, ok, err := resolveTenantChatModelByInstanceID(tid, modelRef, apiKey, baseURL) + if err != nil { + return "", "", apiKey, baseURL, false, err + } + if ok { + return modelName, provider, instanceKey, instanceBaseURL, true, nil + } + 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) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return "", "", apiKey, baseURL, false, nil + } + return "", "", apiKey, baseURL, false, fmt.Errorf("resolve tenant model id %q: %w", modelID, err) + } + if !strings.EqualFold(strings.TrimSpace(model.Status), "active") { + return "", "", apiKey, baseURL, false, fmt.Errorf("tenant model id %s is disabled", modelID) + } + if !entity.ModelType(model.ModelType).Has(entity.ModelTypeChat) { + return "", "", apiKey, baseURL, false, fmt.Errorf("tenant model id %s cannot be used as chat model", modelID) + } + + provider, err := dao.NewTenantModelProviderDAO().GetByID(model.ProviderID) + if err != nil { + return "", "", apiKey, baseURL, false, fmt.Errorf("resolve provider for tenant model id %q: %w", modelID, err) + } + if provider.TenantID != tid { + return "", "", apiKey, baseURL, false, fmt.Errorf("tenant %s has no access to model id %s", tid, modelID) + } + + instance, err := dao.NewTenantModelInstanceDAO().GetByID(model.InstanceID) + if err != nil { + return "", "", apiKey, baseURL, false, fmt.Errorf("resolve instance for tenant model id %q: %w", modelID, err) + } + apiKey, baseURL = fillInstanceCredentials(instance, apiKey, baseURL) + 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) + 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) + if err != nil { + return "", "", apiKey, baseURL, false, fmt.Errorf("resolve provider for tenant model instance id %q: %w", instanceID, err) + } + if provider.TenantID != tid { + 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) + if err != nil { + return "", "", apiKey, baseURL, false, fmt.Errorf("resolve models for tenant model instance id %q: %w", instanceID, err) + } + candidates := make([]*entity.TenantModel, 0, len(models)) + for _, model := range models { + if model == nil || model.ProviderID != provider.ID { + continue + } + if !strings.EqualFold(strings.TrimSpace(model.Status), "active") { + continue + } + if !entity.ModelType(model.ModelType).Has(entity.ModelTypeChat) { + continue + } + candidates = append(candidates, model) + } + if len(candidates) == 0 { + return "", "", apiKey, baseURL, false, fmt.Errorf("tenant model instance id %s has no active chat model", instanceID) + } + if len(candidates) > 1 { + return "", "", apiKey, baseURL, false, fmt.Errorf("tenant model instance id %s has %d active chat models; use tenant_model.id or model@instance@provider", instanceID, len(candidates)) + } + apiKey, baseURL = fillInstanceCredentials(instance, apiKey, baseURL) + return candidates[0].ModelName, provider.ProviderName, apiKey, baseURL, true, nil +} + +func fillInstanceCredentials(instance *entity.TenantModelInstance, apiKey, baseURL string) (string, string) { + if instance == nil { + return apiKey, baseURL + } + if apiKey == "" { + apiKey = instance.APIKey + } + if baseURL == "" && instance.Extra != "" { + var extra map[string]string + if err := json.Unmarshal([]byte(instance.Extra), &extra); err == nil { + baseURL = extra["base_url"] + } + } + return apiKey, baseURL +} + +func isBareTenantModelID(s string) bool { + s = strings.TrimSpace(s) + if len(s) != 32 || strings.Contains(s, "@") { + return false + } + for _, r := range s { + switch { + case r >= '0' && r <= '9': + case r >= 'a' && r <= 'f': + case r >= 'A' && r <= 'F': + default: + return false + } + } + return true +} + // parseLLMIDParts splits a composite llm_id into model, instance, and // provider segments. // diff --git a/internal/handler/agent_component.go b/internal/handler/agent_component.go index b7c6cbf801..d3115221e5 100644 --- a/internal/handler/agent_component.go +++ b/internal/handler/agent_component.go @@ -217,7 +217,9 @@ func (h *AgentHandler) DebugComponent(c *gin.Context) { } invokeCtx := runtime.WithState(c.Request.Context(), debugState) - outputs, err := comp.Invoke(invokeCtx, inputs) + outputs, err := runtime.TrackElapsed(name, func() (map[string]any, error) { + return comp.Invoke(invokeCtx, inputs) + }) if err != nil { common.ResponseWithCodeData(c, common.CodeServerError, nil, "invoke: "+err.Error()) return