From 466f33e6b44f48bb5bd41a77388536c878429773 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Mon, 13 Jul 2026 16:57:56 +0800 Subject: [PATCH] fix: support component brwoser (#16848) ### Summary As title Unable to test it since I don't have apiKey for `openai` , `Anthropic` and `Gemini` --- image --- internal/agent/component/browser.go | 203 ++++++++++++++++-- internal/agent/component/browser_test.go | 162 +++++++++++++- internal/agent/component/categorize_test.go | 1 + internal/agent/component/stagehand_runtime.go | 45 +++- .../agent/component/stagehand_runtime_test.go | 33 +++ 5 files changed, 403 insertions(+), 41 deletions(-) diff --git a/internal/agent/component/browser.go b/internal/agent/component/browser.go index c535dc6b7f..7d7372aed9 100644 --- a/internal/agent/component/browser.go +++ b/internal/agent/component/browser.go @@ -44,16 +44,56 @@ import ( "encoding/json" "errors" "fmt" + "net/url" "regexp" "strings" "ragflow/internal/agent/runtime" "ragflow/internal/dao" "ragflow/internal/entity" + + "gorm.io/gorm" ) const componentNameBrowser = "Browser" +var stagehandNativeProviders = map[string]string{ + "anthropic": "anthropic", + "google": "google", + "openai": "openai", +} + +var browserFactoryDefaultBaseURL = map[string]string{ + "302.ai": "https://api.302.ai/v1", + "anthropic": "https://api.anthropic.com/", + "astraflow": "https://api-us-ca.umodelverse.ai/v1", + "astraflow-cn": "https://api.modelverse.cn/v1", + "avian": "https://api.avian.io/v1", + "cometapi": "https://api.cometapi.com/v1", + "dashscope": "https://dashscope.aliyuncs.com/compatible-mode/v1", + "deepseek": "https://api.deepseek.com/v1", + "deerapi": "https://api.deerapi.com/v1", + "futurmix": "https://futurmix.ai/v1", + "giteeai": "https://ai.gitee.com/v1/", + "hunyuan": "https://api.hunyuan.cloud.tencent.com/v1", + "jiekouai": "https://api.jiekou.ai/openai", + "lingyi-ai": "https://api.lingyiwanwu.com/v1", + "longcat": "https://api.longcat.chat/openai", + "minimax": "https://api.minimaxi.com/v1", + "moonshot": "https://api.moonshot.cn/v1", + "n1n": "https://api.n1n.ai/v1", + "novitaai": "https://api.novita.ai/v3/openai", + "openai": "https://api.openai.com/v1", + "openrouter": "https://openrouter.ai/api/v1", + "perfxcloud": "https://cloud.perfxlab.cn/v1", + "ppio": "https://api.ppinfra.com/v3/openai", + "siliconflow": "https://api.siliconflow.cn/v1", + "stepfun": "https://api.stepfun.com/v1", + "tongyi-qianwen": "https://dashscope.aliyuncs.com/compatible-mode/v1", + "upstage": "https://api.upstage.ai/v1/solar", + "zhipu-ai": "https://open.bigmodel.cn/api/paas/v4", +} + // browserParam is the static DSL param surface for the Browser // component. Mirrors Python `browser.py:LLMParam + browser knobs`: // @@ -228,10 +268,10 @@ func (b *BrowserComponent) Name() string { return b.name } // StagehandInvoker.RunExtract with a `{"type": "string"}` schema. // The flow: // -// 1. Pull tenant_id from `state.Sys["user_id"]`. +// 1. Pull tenant_id from `state.Sys["tenant_id"]`. // 2. Resolve the `prompts` template via `runtime.ResolveTemplate`. -// 3. Split `llm_id` → `(modelName, factory)` and look up the -// tenant LLM config (apiKey, baseURL) from the DAO. +// 3. Resolve `llm_id` as a tenant_model id or legacy model@factory +// reference and look up the tenant LLM config (apiKey, baseURL). // 4. Build `RunExtractRequest` with `ModelName = "openai/"`, // the resolved apiKey/baseURL/instruction, and // `Schema = {"type": "string"}`. @@ -251,29 +291,33 @@ func (b *BrowserComponent) Invoke(ctx context.Context, inputs map[string]any) (m return nil, errors.New("Browser: nil canvas state") } - tenantID, _ := state.Sys["user_id"].(string) + tenantID, _ := state.Sys["tenant_id"].(string) if tenantID == "" { - return nil, errors.New("Browser: tenant_id missing from canvas state (state.Sys[\"user_id\"])") + return nil, errors.New("Browser: tenant_id missing from canvas state (state.Sys[\"tenant_id\"])") } // 1. Resolve prompts template. - resolvedPrompts, err := runtime.ResolveTemplate(b.param.Prompts, state) + prompts := b.param.Prompts + if v, ok := inputs["prompts"].(string); ok && strings.TrimSpace(v) != "" { + prompts = v + } + resolvedPrompts, err := runtime.ResolveTemplate(prompts, state) if err != nil { return nil, fmt.Errorf("Browser: resolve prompts template: %w", err) } // 2. Look up tenant model config. - modelName, factory := resolveLLMID(b.param.LLMID) - apiKey, baseURL, err := resolveTenantLLM(tenantID, modelName, factory) + providerName, modelName, apiKey, baseURL, err := resolveBrowserLLM(tenantID, b.param.LLMID) if err != nil { - return nil, fmt.Errorf("Browser: tenant llm lookup (%q, factory=%q): %w", b.param.LLMID, factory, err) + return nil, fmt.Errorf("Browser: tenant llm lookup (%q): %w", b.param.LLMID, err) } + baseURL = strings.TrimSpace(baseURL) // 3. Build RunExtractRequest with single-string schema. req := RunExtractRequest{ TenantID: tenantID, LLMID: b.param.LLMID, - ModelName: "openai/" + modelName, + ModelName: stagehandModelName(providerName, modelName), BaseURL: baseURL, APIKey: apiKey, Headless: b.param.Headless, @@ -285,7 +329,8 @@ func (b *BrowserComponent) Invoke(ctx context.Context, inputs map[string]any) (m invoker := getDefaultStagehandInvoker() rawJSON, err := invoker.RunExtract(ctx, req) if err != nil { - return nil, fmt.Errorf("Browser: stagehand extract: %w", err) + 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. @@ -302,11 +347,43 @@ func (b *BrowserComponent) Invoke(ctx context.Context, inputs map[string]any) (m "status": 0, "size": len(content), "model_id": b.param.LLMID, - "prompt": b.param.Prompts, + "prompt": prompts, } return out, nil } +func browserBaseURLForLog(raw string) string { + raw = strings.TrimSpace(raw) + if raw == "" { + return "" + } + u, err := url.Parse(raw) + if err != nil { + return "" + } + u.User = nil + return u.String() +} + +func stagehandModelName(providerName, modelName string) string { + modelName = strings.TrimSpace(modelName) + if modelName == "" { + return "" + } + if strings.Contains(modelName, "/") { + prefix, _, _ := strings.Cut(modelName, "/") + if stagehandNativeProviders[strings.ToLower(strings.TrimSpace(prefix))] != "" { + return modelName + } + return "openai/" + modelName + } + providerName = strings.ToLower(strings.TrimSpace(providerName)) + if nativeProvider := stagehandNativeProviders[providerName]; nativeProvider != "" { + return nativeProvider + "/" + modelName + } + return "openai/" + modelName +} + // Stream mirrors Invoke; Browser is a single-shot generator. func (b *BrowserComponent) Stream(ctx context.Context, inputs map[string]any) (<-chan map[string]any, error) { out, err := b.Invoke(ctx, inputs) @@ -337,7 +414,7 @@ func (b *BrowserComponent) Inputs() map[string]string { func (b *BrowserComponent) GetInputForm() map[string]any { return map[string]any{ "prompts": map[string]any{ - "type": "text", + "type": "paragraph", "name": "Prompts", }, "upload_sources": map[string]any{ @@ -360,21 +437,107 @@ func (b *BrowserComponent) Outputs() map[string]string { } } -// resolveTenantLLM looks up the tenant LLM config and returns -// (apiKey, baseURL, modelName). baseURL may be empty when the -// tenant's provider doesn't configure a custom endpoint (the -// stagehand server will then use its openai-compat default). +// resolveBrowserLLM resolves the Browser's selected model into the model name +// and credentials required by the stagehand runtime. It first tries a +// tenant_model.id lookup, then model@factory parsing via resolveTenantLLM. // // 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) { + if tenantLLMLookupForTest != nil { + oldModelName, factory := resolveLLMID(llmID) + apiKey, baseURL, err = tenantLLMLookupForTest(tenantID, oldModelName, factory) + baseURL = browserOpenAICompatibleBaseURL(baseURL, factory) + return factory, oldModelName, apiKey, baseURL, err + } + + providerName, modelName, apiKey, baseURL, err = resolveTenantModelBrowserLLM(tenantID, llmID) + if err == nil { + baseURL = browserOpenAICompatibleBaseURL(baseURL, providerName) + return providerName, modelName, apiKey, baseURL, nil + } + modelErr := err + + oldModelName, factory := resolveLLMID(llmID) + apiKey, baseURL, oldErr := resolveTenantLLM(tenantID, oldModelName, factory) + if oldErr == nil { + baseURL = browserOpenAICompatibleBaseURL(baseURL, factory) + return factory, oldModelName, apiKey, baseURL, nil + } + 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) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return "", "", "", "", err + } + return "", "", "", "", fmt.Errorf("tenant model id=%s lookup failed: %w", modelID, err) + } + if modelRow.Status != "active" { + return "", "", "", "", fmt.Errorf("tenant model id=%s is disabled", modelID) + } + if !entity.ModelType(modelRow.ModelType).Has(entity.ModelTypeChat) { + return "", "", "", "", fmt.Errorf("tenant model id=%s cannot be used as %s model", modelID, entity.ModelTypeChat.String()) + } + + provider, err := dao.NewTenantModelProviderDAO().GetByID(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) + } + return "", "", "", "", err + } + if provider == nil { + return "", "", "", "", fmt.Errorf("provider id=%s not found for model id=%s", modelRow.ProviderID, modelID) + } + if provider.TenantID != tenantID { + return "", "", "", "", fmt.Errorf("tenant %s has no access to provider owned by tenant %s", tenantID, provider.TenantID) + } + + instance, err := dao.NewTenantModelInstanceDAO().GetByID(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) + } + return "", "", "", "", err + } + if instance == nil { + return "", "", "", "", fmt.Errorf("instance id=%s not found for model id=%s", modelRow.InstanceID, modelID) + } + + apiKey = instance.APIKey + if strings.TrimSpace(instance.Extra) != "" { + var extra map[string]string + if err := json.Unmarshal([]byte(instance.Extra), &extra); err != nil { + return "", "", "", "", err + } + baseURL = extra["base_url"] + } + return provider.ProviderName, modelRow.ModelName, apiKey, baseURL, nil +} + +func browserOpenAICompatibleBaseURL(baseURL, provider string) string { + baseURL = strings.TrimSpace(baseURL) + if baseURL != "" { + return baseURL + } + provider = strings.TrimSpace(provider) + if provider == "" { + return "" + } + return browserFactoryDefaultBaseURL[strings.ToLower(provider)] +} + +// resolveTenantLLM looks up the legacy tenant_llm config and returns +// (apiKey, baseURL). baseURL may be empty when the tenant's provider doesn't +// configure a custom endpoint. // // 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) { - if tenantLLMLookupForTest != nil { - return tenantLLMLookupForTest(tenantID, modelName, factory) - } dao := dao.NewTenantLLMDAO() var ( row *entity.TenantLLM diff --git a/internal/agent/component/browser_test.go b/internal/agent/component/browser_test.go index a1c89361e7..ebc9bc6125 100644 --- a/internal/agent/component/browser_test.go +++ b/internal/agent/component/browser_test.go @@ -25,6 +25,7 @@ import ( "ragflow/internal/agent/canvas" "ragflow/internal/agent/runtime" + "ragflow/internal/entity" ) // mockStagehandInvoker captures RunExtract requests and returns a @@ -145,8 +146,8 @@ func TestBrowser_ResolvesSysQueryTemplate(t *testing.T) { "prompts": "{sys.query}打开百度，搜索'2026年最新AI技术趋势'", }) ctx := stateWith(t, map[string]any{ - "user_id": "tenant-1", - "query": "what's the latest", + "tenant_id": "tenant-1", + "query": "what's the latest", }) // LLM lookup will fail (no DB seeded in test), so we can't @@ -196,7 +197,7 @@ func TestBrowser_DispatchesToRuntime(t *testing.T) { "llm_id": "deepseek-v4-pro@DeepSeek", "prompts": "do something", }) - ctx := stateWith(t, map[string]any{"user_id": "tenant-1"}) + ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) _, err := c.Invoke(ctx, nil) if err == nil { @@ -213,8 +214,60 @@ func TestBrowser_DispatchesToRuntime(t *testing.T) { } } -// TestBrowser_MissingTenant: state.Sys["user_id"] is the only -// tenant handle (until the cross-cutting tenant_id fix lands). +func TestResolveBrowserLLM_ResolvesTenantModelID(t *testing.T) { + db := setupComponentTestDB(t) + pushComponentDB(t, db) + if err := db.Create(&entity.TenantModelProvider{ + ID: "provider-1", + TenantID: "tenant-1", + ProviderName: "OpenAI", + }).Error; err != nil { + t.Fatalf("create provider: %v", err) + } + if err := db.Create(&entity.TenantModelInstance{ + ID: "instance-1", + ProviderID: "provider-1", + InstanceName: "default", + APIKey: "sk-instance", + Status: "active", + Extra: `{"base_url":"https://instance.example/v1"}`, + }).Error; err != nil { + t.Fatalf("create instance: %v", err) + } + if err := db.Create(&entity.TenantModel{ + ID: "tenant-model-1", + ProviderID: "provider-1", + InstanceID: "instance-1", + ModelName: "gpt-4o", + ModelType: int(entity.ModelTypeChat), + Status: "active", + }).Error; err != nil { + t.Fatalf("create model: %v", err) + } + + prevLookup := tenantLLMLookupForTest + tenantLLMLookupForTest = nil + t.Cleanup(func() { tenantLLMLookupForTest = prevLookup }) + + provider, model, apiKey, baseURL, err := resolveBrowserLLM("tenant-1", "tenant-model-1") + if err != nil { + t.Fatalf("resolveBrowserLLM: %v", err) + } + if provider != "OpenAI" { + t.Fatalf("provider=%q, want OpenAI", provider) + } + if model != "gpt-4o" { + t.Fatalf("model=%q, want gpt-4o", model) + } + if apiKey != "sk-instance" { + t.Fatalf("apiKey=%q, want sk-instance", apiKey) + } + if baseURL != "https://instance.example/v1" { + t.Fatalf("baseURL=%q, want https://instance.example/v1", baseURL) + } +} + +// TestBrowser_MissingTenant: state.Sys["tenant_id"] is required. // Missing tenant_id must error. func TestBrowser_MissingTenant(t *testing.T) { mock := &mockStagehandInvoker{rawJSON: `"ok"`} @@ -224,7 +277,7 @@ func TestBrowser_MissingTenant(t *testing.T) { "llm_id": "gpt-4o@OpenAI", "prompts": "x", }) - // state with no user_id + // state with no tenant_id state := canvas.NewCanvasState("run-1", "task-1") ctx := canvas.WithState(context.Background(), state) @@ -235,7 +288,7 @@ func TestBrowser_MissingTenant(t *testing.T) { } // TestBrowser_PropagatesRuntimeError: a runtime error surfaces -// wrapped as `Browser: stagehand run: ...`. +// wrapped with model context. // // We can't easily seed the tenant DAO in a unit test, so this test // verifies the error-wrapping contract by injecting a mock runtime @@ -258,7 +311,7 @@ func TestBrowser_PropagatesRuntimeError(t *testing.T) { "llm_id": "gpt-4o@OpenAI", "prompts": "x", }) - ctx := stateWith(t, map[string]any{"user_id": "tenant-1"}) + ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) _, err := c.Invoke(ctx, nil) if err == nil { @@ -267,11 +320,98 @@ func TestBrowser_PropagatesRuntimeError(t *testing.T) { if !strings.Contains(err.Error(), "stagehand extract") { t.Errorf("error should mention stagehand extract; got %v", err) } + if !strings.Contains(err.Error(), `model="openai/gpt-4o"`) { + t.Errorf("error should include model context; got %v", err) + } + if !strings.Contains(err.Error(), "base_url=https://api.openai.com/v1") { + t.Errorf("error should include base_url context; got %v", err) + } if !strings.Contains(err.Error(), "boom") { t.Errorf("error should include underlying message; got %v", err) } } +func TestBrowserBaseURLForLog(t *testing.T) { + cases := []struct { + name string + in string + want string + }{ + {name: "empty", in: "", want: ""}, + {name: "blank", in: " \t\n ", want: ""}, + {name: "redacts_userinfo", in: "https://user:pass@example.com/v1", want: "https://example.com/v1"}, + {name: "invalid", in: "://bad", want: ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := browserBaseURLForLog(tc.in); got != tc.want { + t.Fatalf("browserBaseURLForLog(%q) = %q, want %q", tc.in, got, tc.want) + } + }) + } +} + +func TestBrowserOpenAICompatibleBaseURL(t *testing.T) { + cases := []struct { + name string + baseURL string + provider string + want string + }{ + { + name: "explicit_wins", + baseURL: " https://gateway.example/v1 ", + provider: "DeepSeek", + want: "https://gateway.example/v1", + }, + { + name: "deepseek_default", + provider: "DeepSeek", + want: "https://api.deepseek.com/v1", + }, + { + name: "siliconflow_default", + provider: "SILICONFLOW", + want: "https://api.siliconflow.cn/v1", + }, + { + name: "unknown_provider", + provider: "unknown", + want: "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := browserOpenAICompatibleBaseURL(tc.baseURL, tc.provider); got != tc.want { + t.Fatalf("browserOpenAICompatibleBaseURL(%q, %q) = %q, want %q", tc.baseURL, tc.provider, got, tc.want) + } + }) + } +} + +func TestStagehandModelName(t *testing.T) { + cases := []struct { + name string + provider string + model string + want string + }{ + {name: "deepseek_uses_openai_compatible_provider", provider: "DeepSeek", model: "deepseek-v4-pro", want: "openai/deepseek-v4-pro"}, + {name: "siliconflow_uses_openai_compatible_provider", provider: "SILICONFLOW", model: "Qwen/Qwen3-32B", want: "openai/Qwen/Qwen3-32B"}, + {name: "openai_native", provider: "OpenAI", model: "gpt-4o", want: "openai/gpt-4o"}, + {name: "anthropic_native", provider: "Anthropic", model: "claude-sonnet-4-5", want: "anthropic/claude-sonnet-4-5"}, + {name: "google_native", provider: "Google", model: "gemini-3-flash-preview", want: "google/gemini-3-flash-preview"}, + {name: "already_prefixed", provider: "DeepSeek", model: "openai/deepseek-v4-pro", want: "openai/deepseek-v4-pro"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := stagehandModelName(tc.provider, tc.model); got != tc.want { + t.Fatalf("stagehandModelName(%q, %q) = %q, want %q", tc.provider, tc.model, got, tc.want) + } + }) + } +} + // TestBrowser_RunExtractRequestShape verifies the RunExtractRequest // fields forwarded to the runtime: ModelName, TenantID, APIKey, // Schema, and the resolved Instruction. @@ -289,7 +429,7 @@ func TestBrowser_RunExtractRequestShape(t *testing.T) { "llm_id": "gpt-4o@OpenAI", "prompts": "extract the page title", }) - ctx := stateWith(t, map[string]any{"user_id": "tenant-1"}) + ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) if _, err := c.Invoke(ctx, nil); err != nil { t.Fatalf("Invoke: %v", err) @@ -332,7 +472,7 @@ func TestBrowser_HeadlessPropagates(t *testing.T) { "prompts": "x", "headless": false, }) - ctx := stateWith(t, map[string]any{"user_id": "tenant-1"}) + ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) if _, err := c.Invoke(ctx, nil); err != nil { t.Fatalf("Invoke: %v", err) @@ -363,7 +503,7 @@ func TestBrowser_OutputsShape(t *testing.T) { "llm_id": "gpt-4o@OpenAI", "prompts": "x", }) - ctx := stateWith(t, map[string]any{"user_id": "tenant-1"}) + ctx := stateWith(t, map[string]any{"tenant_id": "tenant-1"}) out, err := c.Invoke(ctx, nil) if err != nil { diff --git a/internal/agent/component/categorize_test.go b/internal/agent/component/categorize_test.go index c71b62e6c0..57f80a00e9 100644 --- a/internal/agent/component/categorize_test.go +++ b/internal/agent/component/categorize_test.go @@ -390,6 +390,7 @@ func setupComponentTestDB(t *testing.T) *gorm.DB { &entity.TenantLLM{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, + &entity.TenantModel{}, ); err != nil { t.Fatalf("migrate: %v", err) } diff --git a/internal/agent/component/stagehand_runtime.go b/internal/agent/component/stagehand_runtime.go index cf9466c2f7..07c01ab360 100644 --- a/internal/agent/component/stagehand_runtime.go +++ b/internal/agent/component/stagehand_runtime.go @@ -67,6 +67,7 @@ import ( "math" "ragflow/internal/common" "strconv" + "strings" "sync" "sync/atomic" "time" @@ -427,12 +428,7 @@ func (r *stagehandRuntime) RunTask(ctx context.Context, req RunTaskRequest) (str // OfSessionExecutesAgentConfigModelGenericModelConfigObject // variant is the only path that exposes BaseURL. Model: stagehand.SessionExecuteParamsAgentConfigModelUnion{ - OfSessionExecutesAgentConfigModelGenericModelConfigObject: &stagehand.SessionExecuteParamsAgentConfigModelGenericModelConfigObject{ - ModelName: req.ModelName, - APIKey: stagehand.String(req.APIKey), - BaseURL: stagehand.String(req.BaseURL), - Provider: "openai", - }, + OfSessionExecutesAgentConfigModelGenericModelConfigObject: stagehandExecuteModelConfig(req), }, }, ExecuteOptions: stagehand.SessionExecuteParamsExecuteOptions{ @@ -455,6 +451,18 @@ func (r *stagehandRuntime) RunTask(ctx context.Context, req RunTaskRequest) (str return execResp.Data.Result.Message, nil } +func stagehandExecuteModelConfig(req RunTaskRequest) *stagehand.SessionExecuteParamsAgentConfigModelGenericModelConfigObject { + model := &stagehand.SessionExecuteParamsAgentConfigModelGenericModelConfigObject{ + ModelName: req.ModelName, + APIKey: stagehand.String(req.APIKey), + Provider: "openai", + } + if baseURL := strings.TrimSpace(req.BaseURL); baseURL != "" { + model.BaseURL = stagehand.String(baseURL) + } + return model +} + // RunExtractRequest is the input to stagehandRuntime.RunExtract. // // RunExtract is the "structured single-shot" sibling of RunTask: @@ -493,9 +501,9 @@ type RunExtractRequest struct { // 2. Sessions.Start (browser, model, launch options). // 3. defer Sessions.End. // 4. If URL != "", Sessions.Navigate({URL: req.URL}). -// 5. Sessions.Extract({Instruction, Schema}). The response -// `Data.Result` field is the structured data matching Schema; -// we marshal it as JSON and return. +// 5. Sessions.Extract({Instruction, Schema, Options.Model}). The +// response `Data.Result` field is the structured data matching +// Schema; we marshal it as JSON and return. // // RunExtract reuses the same `stagehandRuntime` cache as RunTask // (per `(apiKey, baseURL, modelName)` key), so cold-start cost is @@ -564,7 +572,10 @@ func (r *stagehandRuntime) RunExtract(ctx context.Context, req RunExtractRequest extractResp, err := client.Sessions.Extract(ctx, sessionID, stagehand.SessionExtractParams{ Instruction: stagehand.String(req.Instruction), - Schema: req.Schema, + Options: stagehand.SessionExtractParamsOptions{ + Model: stagehandExtractModelConfig(req), + }, + Schema: req.Schema, }) if err != nil { return "", fmt.Errorf("stagehand runtime: RunExtract: Sessions.Extract: %w", err) @@ -583,6 +594,20 @@ func (r *stagehandRuntime) RunExtract(ctx context.Context, req RunExtractRequest return string(out), nil } +func stagehandExtractModelConfig(req RunExtractRequest) stagehand.SessionExtractParamsOptionsModelUnion { + model := &stagehand.SessionExtractParamsOptionsModelGenericModelConfigObject{ + ModelName: req.ModelName, + APIKey: stagehand.String(req.APIKey), + Provider: "openai", + } + if baseURL := strings.TrimSpace(req.BaseURL); baseURL != "" { + model.BaseURL = stagehand.String(baseURL) + } + return stagehand.SessionExtractParamsOptionsModelUnion{ + OfSessionExtractsOptionsModelGenericModelConfigObject: model, + } +} + // Close drains the sweeper goroutine and closes every cached // client (each Close kills its stagehand-server subprocess). // Safe to call multiple times. diff --git a/internal/agent/component/stagehand_runtime_test.go b/internal/agent/component/stagehand_runtime_test.go index 3a237fa7e4..b6943f25a0 100644 --- a/internal/agent/component/stagehand_runtime_test.go +++ b/internal/agent/component/stagehand_runtime_test.go @@ -69,6 +69,39 @@ func TestStagehandRuntime_ValidatesRequiredFields(t *testing.T) { } } +func TestStagehandRuntime_RunExtractModelConfig(t *testing.T) { + model := stagehandExtractModelConfig(RunExtractRequest{ + ModelName: "openai/gpt-4o", + BaseURL: "https://example.test/v1", + APIKey: "sk-test", + }) + + if got := model.GetModelName(); got == nil || *got != "openai/gpt-4o" { + t.Fatalf("modelName = %v, want openai/gpt-4o", got) + } + if got := model.GetAPIKey(); got == nil || *got != "sk-test" { + t.Fatalf("apiKey = %v, want sk-test", got) + } + if got := model.GetBaseURL(); got == nil || *got != "https://example.test/v1" { + t.Fatalf("baseURL = %v, want https://example.test/v1", got) + } + if got := model.GetProvider(); got == nil || *got != "openai" { + t.Fatalf("provider = %v, want openai", got) + } +} + +func TestStagehandRuntime_RunExtractModelConfigOmitsEmptyBaseURL(t *testing.T) { + model := stagehandExtractModelConfig(RunExtractRequest{ + ModelName: "openai/gpt-4o", + BaseURL: " \t\n ", + APIKey: "sk-test", + }) + + if got := model.GetBaseURL(); got != nil { + t.Fatalf("baseURL = %q, want omitted", *got) + } +} + // --------------------------------------------------------------------------- // Cache hit / miss / key formula // ---------------------------------------------------------------------------