Fix developer agent broken when the default OLAP connector is invalid (#8765)

Author: begelundmuller

runtime/ai/develop_file.go

@@ -9,7 +9,6 @@ import (
 	aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1"
 	"github.com/rilldata/rill/runtime"
 	"github.com/rilldata/rill/runtime/ai/instructions"
-	"github.com/rilldata/rill/runtime/drivers"
 )
 
 const DevelopFileName = "develop_file"
@@ -137,37 +136,21 @@ func (t *DevelopFile) Handler(ctx context.Context, args *DevelopFileArgs) (*Deve
 
 func (t *DevelopFile) userPrompt(ctx context.Context, args *DevelopFileArgs) (string, error) {
 	// Get default OLAP info
-	olapConnector, olapDriver, olapModeReadwrite, err := defaultOLAPInfo(ctx, t.Runtime, GetSession(ctx).InstanceID())
+	olapInfo, err := defaultOLAPInfo(ctx, t.Runtime, GetSession(ctx).InstanceID())
 	if err != nil {
 		return "", err
 	}
 
 	// Prepare template data.
 	session := GetSession(ctx)
 	data := map[string]any{
-		"path":                   args.Path,
-		"type":                   args.Type,
-		"prompt":                 args.Prompt,
-		"ai_instructions":        session.ProjectInstructions(),
-		"default_olap_connector": olapConnector,
-		"default_olap_driver":    olapDriver,
-		"default_olap_readwrite": olapModeReadwrite,
+		"path":              args.Path,
+		"type":              args.Type,
+		"prompt":            args.Prompt,
+		"ai_instructions":   session.ProjectInstructions(),
+		"default_olap_info": olapInfo,
 	}
 
-	// Add OLAP dialect
-	olap, release, err := t.Runtime.OLAP(ctx, session.InstanceID(), "")
-	if err != nil {
-		return "", err
-	}
-	defer release()
-	dialect := olap.Dialect()
-	if dialect == drivers.DialectUnspecified {
-		dialect = drivers.DialectDuckDB
-	}
-	data["dialect"] = dialect.String()
-
-	// - Additional prompt: Don’t do much discovery, stay aligned on your specific task. Failure is okay, your parent will re-evaluate.
-
 	// Generate the user prompt
 	return executeTemplate(`
 You should develop a Rill project file based on the following task description:
@@ -180,8 +163,7 @@ Here is some important context:
 - When you call 'write_file', if it returns a parse or reconcile error, do your best to fix the issue and try again. If you think the error is unrelated to the current path, let the parent agent know to handle it.
 
 Here is some additional context that may or may not be relevant to your task:
-- The project's default OLAP connector is named {{ .default_olap_connector }} (driver: {{ .default_olap_driver }}).
-{{ if .ai_instructions }}- The user has configured global additional instructions for you. They may not relate to the current request, and may not even relate to your work as a data engineer agent. Only use them if you find them relevant. They are: {{ .ai_instructions }}
-{{ end }}
+- Info about the project's default OLAP connector: {{ .default_olap_info }}.
+{{ if .ai_instructions }}- The user has configured global additional instructions for you. They may not relate to the current request, and may not even relate to your work as a data engineer agent. Only use them if you find them relevant. They are: {{ .ai_instructions }}{{ end }}
 `, data)
 }

runtime/ai/developer_agent.go

@@ -2,6 +2,7 @@ package ai
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strings"
 
@@ -121,21 +122,19 @@ func (t *DeveloperAgent) systemPrompt() (string, error) {
 
 func (t *DeveloperAgent) userPrompt(ctx context.Context, args *DeveloperAgentArgs) (string, error) {
 	// Get default OLAP info
-	olapConnector, olapDriver, olapModeReadwrite, err := defaultOLAPInfo(ctx, t.Runtime, GetSession(ctx).InstanceID())
+	olapInfo, err := defaultOLAPInfo(ctx, t.Runtime, GetSession(ctx).InstanceID())
 	if err != nil {
 		return "", err
 	}
 
 	// Prepare template data.
 	session := GetSession(ctx)
 	data := map[string]any{
-		"prompt":                 args.Prompt,
-		"init_project":           args.InitProject,
-		"current_file_path":      args.CurrentFilePath,
-		"ai_instructions":        session.ProjectInstructions(),
-		"default_olap_connector": olapConnector,
-		"default_olap_driver":    olapDriver,
-		"default_olap_readwrite": olapModeReadwrite,
+		"prompt":            args.Prompt,
+		"init_project":      args.InitProject,
+		"current_file_path": args.CurrentFilePath,
+		"ai_instructions":   session.ProjectInstructions(),
+		"default_olap_info": olapInfo,
 	}
 
 	// Generate the system prompt
@@ -157,9 +156,8 @@ Feel free to change the default OLAP connector if it makes sense for the task.
 The user has configured global additional instructions for you. They may not relate to the current request, and may not even relate to your work as a data engineer agent. Only use them if you find them relevant. They are: {{ .ai_instructions }}
 {{ end }}
 
-This may not relate to the user's task, but for context, the project's default OLAP connector is named {{ .default_olap_connector }} (driver: {{ .default_olap_driver }}).
-{{ if .default_olap_readwrite }} The default OLAP is in readwrite mode, so you can use it in models if you want.
-{{ else }} The default OLAP is in read-only mode, so you cannot create models in it. {{ end }}
+For context, here are some details about the project's default OLAP connector: {{ .default_olap_info }}.
+Note that you can only use it in model resources if it is not readonly.
 
 Call "navigate" tool for the main file created/edited in the conversation. Use kind "file" and pass the written file path.
 Prefer dashboard or metrics view files over other files.
@@ -192,25 +190,37 @@ func checkDeveloperAccess(ctx context.Context, rt *runtime.Runtime, internal boo
 	return ff["developer_agent"], nil
 }
 
-// defaultOLAPInfo retrieves the default OLAP connector name, driver, and whether it is in readwrite mode.
-func defaultOLAPInfo(ctx context.Context, rt *runtime.Runtime, instanceID string) (string, string, bool, error) {
+// defaultOLAPInfo returns info about the default OLAP connector name, driver, and whether it is in readwrite mode.
+func defaultOLAPInfo(ctx context.Context, rt *runtime.Runtime, instanceID string) (string, error) {
 	// Get instance config
 	inst, err := rt.Instance(ctx, instanceID)
 	if err != nil {
-		return "", "", false, err
+		return "", err
 	}
 	olapConnector := inst.ResolveOLAPConnector()
 
 	// Open OLAP handle
-	olap, release, err := rt.AcquireHandle(ctx, instanceID, olapConnector)
+	handle, release, err := rt.AcquireHandle(ctx, instanceID, olapConnector)
 	if err != nil {
-		return "", "", false, err
+		if errors.Is(err, ctx.Err()) {
+			return "", err
+		}
+		return fmt.Sprintf("name: %q, error: %v", olapConnector, err), nil
 	}
 	defer release()
 
+	// Get dialect
+
+	// Add OLAP dialect
+	olap, ok := handle.AsOLAP(instanceID)
+	if !ok {
+		return fmt.Sprintf("name: %q, error: not an OLAP connector", olapConnector), nil
+	}
+	dialect := olap.Dialect()
+
 	// Determine if OLAP can run models
-	_, err = olap.AsModelManager(instanceID)
+	_, err = handle.AsModelManager(instanceID)
 	olapModeReadwrite := err == nil
 
-	return olapConnector, olap.Driver(), olapModeReadwrite, nil
+	return fmt.Sprintf("name: %q, driver: %q, dialect: %q, readonly: %v", olapConnector, handle.Driver(), dialect.String(), !olapModeReadwrite), nil
 }

runtime/ai/developer_agent_test.go

@@ -253,3 +253,38 @@ measures:
 	// Check that it doesn't have any parse or reconcile errors.
 	testruntime.RequireReconcileState(t, rt, instanceID, 5, 0, 0)
 }
+
+func TestInvalidDefaultOLAP(t *testing.T) {
+	// Setup a project with a misconfigured default OLAP
+	rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{
+		AIConnector: "openai",
+		Files: map[string]string{
+			"rill.yaml": `olap_connector: duckdb_missing`,
+			"connectors/duckdb.yaml": `
+type: connector
+driver: duckdb
+managed: true
+`,
+		},
+	})
+	testruntime.RequireReconcileState(t, rt, instanceID, 2, 0, 0)
+
+	// Initialize eval
+
+	// Ask it to fix the default OLAP connector
+	s := newEval(t, rt, instanceID)
+	var res *ai.RouterAgentResult
+	_, err := s.CallTool(t.Context(), ai.RoleUser, ai.RouterAgentName, &res, ai.RouterAgentArgs{
+		Prompt: "The default OLAP connector is misconfigured. Can you fix it?",
+	})
+	require.NoError(t, err)
+	require.Equal(t, ai.DeveloperAgentName, res.Agent)
+
+	// Check that it doesn't have any parse or reconcile errors.
+	testruntime.RequireReconcileState(t, rt, instanceID, 2, 0, 0)
+
+	// Check it fixed the default OLAP
+	inst, err := rt.Instance(t.Context(), instanceID)
+	require.NoError(t, err)
+	require.Equal(t, "duckdb", inst.ResolveOLAPConnector())
+}