pkg/aflow: make LLM model per-agent rather than per-flow

Having LLM model per-agent is even more flexible than per-flow.
We can have some more complex tasks during patch generation with the most elaborate model,
but also some simpler ones with less elaborate models.

11 files changed, 29 insertions, 30 deletions

diff --git a/dashboard/app/ai.go b/dashboard/app/ai.go
index 8d8767832..39f1479a3 100644
--- a/dashboard/app/ai.go
+++ b/dashboard/app/ai.go
@@ -48,7 +48,6 @@ type uiAIJob struct {
 	Created          time.Time
 	Started          time.Time
 	Finished         time.Time
-	LLMModel         string
 	CodeRevision     string
 	CodeRevisionLink string
 	Error            string
@@ -68,6 +67,7 @@ type uiAITrajectorySpan struct {
 	Nesting     int64
 	Type        string
 	Name        string
+	Model       string
 	Duration    time.Duration
 	Error       string
 	Args        string
@@ -198,7 +198,6 @@ func makeUIAIJob(job *aidb.Job) *uiAIJob {
 		Created:          job.Created,
 		Started:          nullTime(job.Started),
 		Finished:         nullTime(job.Finished),
-		LLMModel:         job.LLMModel,
 		CodeRevision:     job.CodeRevision,
 		CodeRevisionLink: vcs.LogLink(vcs.SyzkallerRepo, job.CodeRevision),
 		Error:            job.Error,
@@ -220,6 +219,7 @@ func makeUIAITrajectory(trajetory []*aidb.TrajectorySpan) []*uiAITrajectorySpan
 			Nesting:     span.Nesting,
 			Type:        span.Type,
 			Name:        span.Name,
+			Model:       span.Model,
 			Duration:    duration,
 			Error:       nullString(span.Error),
 			Args:        nullJSON(span.Args),
@@ -238,7 +238,7 @@ func apiAIJobPoll(ctx context.Context, req *dashapi.AIJobPollReq) (any, error) {
 		return nil, fmt.Errorf("invalid request")
 	}
 	for _, flow := range req.Workflows {
-		if flow.Type == "" || flow.Name == "" || flow.LLMModel == "" {
+		if flow.Type == "" || flow.Name == "" {
 			return nil, fmt.Errorf("invalid request")
 		}
 	}
diff --git a/dashboard/app/ai_test.go b/dashboard/app/ai_test.go
index b775b2a89..addf71f5a 100644
--- a/dashboard/app/ai_test.go
+++ b/dashboard/app/ai_test.go
@@ -64,9 +64,9 @@ func TestAIBugWorkflows(t *testing.T) {
 	_, err := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{
 		CodeRevision: prog.GitRevision,
 		Workflows: []dashapi.AIWorkflow{
-			{Type: "patching", Name: "patching", LLMModel: "smarty"},
-			{Type: "patching", Name: "patching-foo", LLMModel: "smarty"},
-			{Type: "patching", Name: "patching-bar", LLMModel: "smarty"},
+			{Type: "patching", Name: "patching"},
+			{Type: "patching", Name: "patching-foo"},
+			{Type: "patching", Name: "patching-bar"},
 		},
 	})
 	require.NoError(t, err)
@@ -77,10 +77,10 @@ func TestAIBugWorkflows(t *testing.T) {
 	_, err = c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{
 		CodeRevision: prog.GitRevision,
 		Workflows: []dashapi.AIWorkflow{
-			{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"},
+			{Type: "patching", Name: "patching"},
+			{Type: "patching", Name: "patching-bar"},
+			{Type: "patching", Name: "patching-baz"},
+			{Type: "assessment-kcsan", Name: "assessment-kcsan"},
 		},
 	})
 	require.NoError(t, err)
@@ -88,11 +88,11 @@ func TestAIBugWorkflows(t *testing.T) {
 	_, err = c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{
 		CodeRevision: prog.GitRevision,
 		Workflows: []dashapi.AIWorkflow{
-			{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"},
+			{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"},
 		},
 	})
 	require.NoError(t, err)
@@ -115,7 +115,7 @@ func TestAIJob(t *testing.T) {
 	resp, err := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{
 		CodeRevision: prog.GitRevision,
 		Workflows: []dashapi.AIWorkflow{
-			{Type: "assessment-kcsan", Name: "assessment-kcsan", LLMModel: "smarty"},
+			{Type: "assessment-kcsan", Name: "assessment-kcsan"},
 		},
 	})
 	require.NoError(t, err)
@@ -134,7 +134,7 @@ func TestAIJob(t *testing.T) {
 	resp2, err2 := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{
 		CodeRevision: prog.GitRevision,
 		Workflows: []dashapi.AIWorkflow{
-			{Type: "assessment-kcsan", Name: "assessment-kcsan", LLMModel: "smarty"},
+			{Type: "assessment-kcsan", Name: "assessment-kcsan"},
 		},
 	})
 	require.NoError(t, err2)
@@ -210,7 +210,7 @@ func TestAIAssessmentKCSAN(t *testing.T) {
 	resp, err := c.aiClient.AIJobPoll(&dashapi.AIJobPollReq{
 		CodeRevision: prog.GitRevision,
 		Workflows: []dashapi.AIWorkflow{
-			{Type: ai.WorkflowAssessmentKCSAN, Name: string(ai.WorkflowAssessmentKCSAN), LLMModel: "smarty"},
+			{Type: ai.WorkflowAssessmentKCSAN, Name: string(ai.WorkflowAssessmentKCSAN)},
 		},
 	})
 	require.NoError(t, err)
diff --git a/dashboard/app/aidb/crud.go b/dashboard/app/aidb/crud.go
index 4f7e93f6a..872a70ace 100644
--- a/dashboard/app/aidb/crud.go
+++ b/dashboard/app/aidb/crud.go
@@ -129,12 +129,6 @@ func StartJob(ctx context.Context, req *dashapi.AIJobPollReq) (*Job, error) {
 			job = jobs[0]
 		}
 		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.CodeRevision = req.CodeRevision
 		mut, err := spanner.InsertOrUpdateStruct("Jobs", job)
 		if err != nil {
@@ -184,6 +178,7 @@ func StoreTrajectorySpan(ctx context.Context, jobID string, span *trajectory.Spa
 		Nesting:     int64(span.Nesting),
 		Type:        string(span.Type),
 		Name:        span.Name,
+		Model:       span.Model,
 		Started:     span.Started,
 		Finished:    toNullTime(span.Finished),
 		Error:       toNullString(span.Error),
diff --git a/dashboard/app/aidb/entities.go b/dashboard/app/aidb/entities.go
index 23df884df..0a3e7b164 100644
--- a/dashboard/app/aidb/entities.go
+++ b/dashboard/app/aidb/entities.go
@@ -28,7 +28,6 @@ type Job struct {
 	Created      time.Time
 	Started      spanner.NullTime
 	Finished     spanner.NullTime
-	LLMModel     string // LLM model used to execute the job, filled when the job is started
 	CodeRevision string // syzkaller revision, filled when the job is started
 	Error        string // for finished jobs
 	Args         spanner.NullJSON
@@ -43,6 +42,7 @@ type TrajectorySpan struct {
 	Nesting     int64
 	Type        string
 	Name        string
+	Model       string
 	Started     time.Time
 	Finished    spanner.NullTime
 	Error       spanner.NullString
diff --git a/dashboard/app/aidb/migrations/3_add_trajectory_model.down.sql b/dashboard/app/aidb/migrations/3_add_trajectory_model.down.sql
new file mode 100644
index 000000000..9c8ee7020
--- /dev/null
+++ b/dashboard/app/aidb/migrations/3_add_trajectory_model.down.sql
@@ -0,0 +1 @@
+ALTER TABLE TrajectorySpans DROP COLUMN Model;
diff --git a/dashboard/app/aidb/migrations/3_add_trajectory_model.up.sql b/dashboard/app/aidb/migrations/3_add_trajectory_model.up.sql
new file mode 100644
index 000000000..c5cd8821d
--- /dev/null
+++ b/dashboard/app/aidb/migrations/3_add_trajectory_model.up.sql
@@ -0,0 +1 @@
+ALTER TABLE TrajectorySpans ADD COLUMN Model STRING(1000);
diff --git a/dashboard/app/aidb/migrations/4_remove_jobs_model.down.sql b/dashboard/app/aidb/migrations/4_remove_jobs_model.down.sql
new file mode 100644
index 000000000..1d9885cb8
--- /dev/null
+++ b/dashboard/app/aidb/migrations/4_remove_jobs_model.down.sql
@@ -0,0 +1 @@
+ALTER TABLE Jobs ADD COLUMN LLMModel STRING(1000);
diff --git a/dashboard/app/aidb/migrations/4_remove_jobs_model.up.sql b/dashboard/app/aidb/migrations/4_remove_jobs_model.up.sql
new file mode 100644
index 000000000..85b4c74e4
--- /dev/null
+++ b/dashboard/app/aidb/migrations/4_remove_jobs_model.up.sql
@@ -0,0 +1 @@
+ALTER TABLE Jobs DROP COLUMN LLMModel;
diff --git a/dashboard/app/templates/ai_job.html b/dashboard/app/templates/ai_job.html
index 8f2526c63..f8f2b82bd 100644
--- a/dashboard/app/templates/ai_job.html
+++ b/dashboard/app/templates/ai_job.html
@@ -65,6 +65,9 @@ Detailed info on a single AI job execution.
 			<td>
 				<details>
 					<summary>{{formatDuration $span.Duration}}</summary>
+					{{if $span.Model}}
+						<b>Model:</b> <div id="ai_details_div"><pre>{{$span.Model}}</pre></div><br>
+					{{end}}
 					{{if $span.Error}}
 						<b>Error:</b> <div id="ai_details_div"><pre>{{$span.Error}}</pre></div><br>
 					{{end}}
diff --git a/dashboard/app/templates/templates.html b/dashboard/app/templates/templates.html
index dee2d1300..d20727a73 100644
--- a/dashboard/app/templates/templates.html
+++ b/dashboard/app/templates/templates.html
@@ -689,7 +689,6 @@ Use of this source code is governed by Apache 2 LICENSE that can be found in the
 		<th><a onclick="return sortTable(this, 'Created', textSort)" href="#">Created</a></th>
 		<th><a onclick="return sortTable(this, 'Started', textSort)" href="#">Started</a></th>
 		<th><a onclick="return sortTable(this, 'Finished', textSort)" href="#">Finished</a></th>
-		<th><a onclick="return sortTable(this, 'Model', textSort)" href="#">Model</a></th>
 		<th><a onclick="return sortTable(this, 'Revision', textSort)" href="#">Revision</a></th>
 		<th><a onclick="return sortTable(this, 'Error', textSort)" href="#">Error</a></th>
 	</tr></thead>
@@ -710,7 +709,6 @@ Use of this source code is governed by Apache 2 LICENSE that can be found in the
 		<td>{{formatTime $job.Created}}</td>
 		<td>{{formatTime $job.Started}}</td>
 		<td>{{formatTime $job.Finished}}</td>
-		<td>{{$job.LLMModel}}</td>
 		<td class="tag">{{link $job.CodeRevisionLink $job.CodeRevision}}</td>
 		<td>{{$job.Error}}</td>
 	</tr>
diff --git a/dashboard/dashapi/ai.go b/dashboard/dashapi/ai.go
index 8134e5744..dfa410402 100644
--- a/dashboard/dashapi/ai.go
+++ b/dashboard/dashapi/ai.go
@@ -14,9 +14,8 @@ type AIJobPollReq struct {
 }
 
 type AIWorkflow struct {
-	Type     ai.WorkflowType
-	Name     string
-	LLMModel string // LLM model that will be used to execute this workflow
+	Type ai.WorkflowType
+	Name string
 }
 
 type AIJobPollResp struct {