From a9d6a79219801d2130df3b1a792c57f0e5428e9f Mon Sep 17 00:00:00 2001 From: Dmitry Vyukov Date: Wed, 14 Jan 2026 11:35:20 +0100 Subject: pkg/aflow: allow to specify model per-flow We may want to use a weaker model for some workflows. Allow to use different models for different workflows. --- dashboard/app/ai.go | 7 ++++++- dashboard/app/ai_test.go | 36 +++++++++++++++--------------------- dashboard/app/aidb/crud.go | 10 ++++++---- dashboard/dashapi/ai.go | 6 +++--- 4 files changed, 30 insertions(+), 29 deletions(-) (limited to 'dashboard') diff --git a/dashboard/app/ai.go b/dashboard/app/ai.go index 29bfc38d1..2912addea 100644 --- a/dashboard/app/ai.go +++ b/dashboard/app/ai.go @@ -207,9 +207,14 @@ func makeUIAITrajectory(trajetory []*aidb.TrajectorySpan) []*uiAITrajectorySpan } func apiAIJobPoll(ctx context.Context, req *dashapi.AIJobPollReq) (any, error) { - if len(req.Workflows) == 0 || req.CodeRevision == "" || req.LLMModel == "" { + if len(req.Workflows) == 0 || req.CodeRevision == "" { return nil, fmt.Errorf("invalid request") } + for _, flow := range req.Workflows { + if flow.Type == "" || flow.Name == "" || flow.LLMModel == "" { + return nil, fmt.Errorf("invalid request") + } + } if err := aidb.UpdateWorkflows(ctx, req.Workflows); err != nil { return nil, fmt.Errorf("failed UpdateWorkflows: %w", err) } diff --git a/dashboard/app/ai_test.go b/dashboard/app/ai_test.go index bd3935b18..b775b2a89 100644 --- a/dashboard/app/ai_test.go +++ b/dashboard/app/ai_test.go @@ -63,11 +63,10 @@ func TestAIBugWorkflows(t *testing.T) { _, err := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{ CodeRevision: prog.GitRevision, - LLMModel: "smarty", Workflows: []dashapi.AIWorkflow{ - {Type: "patching", Name: "patching"}, - {Type: "patching", Name: "patching-foo"}, - {Type: "patching", Name: "patching-bar"}, + {Type: "patching", Name: "patching", LLMModel: "smarty"}, + {Type: "patching", Name: "patching-foo", LLMModel: "smarty"}, + {Type: "patching", Name: "patching-bar", LLMModel: "smarty"}, }, }) require.NoError(t, err) @@ -77,25 +76,23 @@ func TestAIBugWorkflows(t *testing.T) { _, err = c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{ CodeRevision: prog.GitRevision, - LLMModel: "smarty", Workflows: []dashapi.AIWorkflow{ - {Type: "patching", Name: "patching"}, - {Type: "patching", Name: "patching-bar"}, - {Type: "patching", Name: "patching-baz"}, - {Type: "assessment-kcsan", Name: "assessment-kcsan"}, + {Type: "patching", Name: "patching", LLMModel: "smarty"}, + {Type: "patching", Name: "patching-bar", LLMModel: "smarty"}, + {Type: "patching", Name: "patching-baz", LLMModel: "smarty"}, + {Type: "assessment-kcsan", Name: "assessment-kcsan", LLMModel: "smarty"}, }, }) require.NoError(t, err) _, err = c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{ CodeRevision: prog.GitRevision, - LLMModel: "smarty", Workflows: []dashapi.AIWorkflow{ - {Type: "patching", Name: "patching"}, - {Type: "patching", Name: "patching-bar"}, - {Type: "patching", Name: "patching-qux"}, - {Type: "assessment-kcsan", Name: "assessment-kcsan"}, - {Type: "assessment-kcsan", Name: "assessment-kcsan-foo"}, + {Type: "patching", Name: "patching", LLMModel: "smarty"}, + {Type: "patching", Name: "patching-bar", LLMModel: "smarty"}, + {Type: "patching", Name: "patching-qux", LLMModel: "smarty"}, + {Type: "assessment-kcsan", Name: "assessment-kcsan", LLMModel: "smarty"}, + {Type: "assessment-kcsan", Name: "assessment-kcsan-foo", LLMModel: "smarty"}, }, }) require.NoError(t, err) @@ -117,9 +114,8 @@ func TestAIJob(t *testing.T) { resp, err := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{ CodeRevision: prog.GitRevision, - LLMModel: "smarty", Workflows: []dashapi.AIWorkflow{ - {Type: "assessment-kcsan", Name: "assessment-kcsan"}, + {Type: "assessment-kcsan", Name: "assessment-kcsan", LLMModel: "smarty"}, }, }) require.NoError(t, err) @@ -137,9 +133,8 @@ func TestAIJob(t *testing.T) { resp2, err2 := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{ CodeRevision: prog.GitRevision, - LLMModel: "smarty", Workflows: []dashapi.AIWorkflow{ - {Type: "assessment-kcsan", Name: "assessment-kcsan"}, + {Type: "assessment-kcsan", Name: "assessment-kcsan", LLMModel: "smarty"}, }, }) require.NoError(t, err2) @@ -214,9 +209,8 @@ func TestAIAssessmentKCSAN(t *testing.T) { resp, err := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{ CodeRevision: prog.GitRevision, - LLMModel: "smarty", Workflows: []dashapi.AIWorkflow{ - {Type: ai.WorkflowAssessmentKCSAN, Name: string(ai.WorkflowAssessmentKCSAN)}, + {Type: ai.WorkflowAssessmentKCSAN, Name: string(ai.WorkflowAssessmentKCSAN), LLMModel: "smarty"}, }, }) require.NoError(t, err) diff --git a/dashboard/app/aidb/crud.go b/dashboard/app/aidb/crud.go index a01c370b6..4f7e93f6a 100644 --- a/dashboard/app/aidb/crud.go +++ b/dashboard/app/aidb/crud.go @@ -128,11 +128,13 @@ func StartJob(ctx context.Context, req *dashapi.AIJobPollReq) (*Job, error) { } job = jobs[0] } - job.Started = spanner.NullTime{ - Time: TimeNow(ctx), - Valid: true, + job.Started = spanner.NullTime{Time: TimeNow(ctx), Valid: true} + for _, flow := range req.Workflows { + if job.Workflow == flow.Name { + job.LLMModel = flow.LLMModel + break + } } - job.LLMModel = req.LLMModel job.CodeRevision = req.CodeRevision mut, err := spanner.InsertOrUpdateStruct("Jobs", job) if err != nil { diff --git a/dashboard/dashapi/ai.go b/dashboard/dashapi/ai.go index 893ed7215..8134e5744 100644 --- a/dashboard/dashapi/ai.go +++ b/dashboard/dashapi/ai.go @@ -9,14 +9,14 @@ import ( ) type AIJobPollReq struct { - LLMModel string // LLM model that will be used to execute jobs CodeRevision string // git commit of the syz-agent server Workflows []AIWorkflow } type AIWorkflow struct { - Type ai.WorkflowType - Name string + Type ai.WorkflowType + Name string + LLMModel string // LLM model that will be used to execute this workflow } type AIJobPollResp struct { -- cgit mrf-deployment