Add support for user system prompt (#217)

* Add support for user system prompt

* Remove accidentally added code

* Rename `systemPrompt` to `user_system_prompt`

Author: markmiro

packages/ai/mod.ts

@@ -136,6 +136,7 @@ function createSystemPrompt(
   currentCellId?: string,
   filepaths?: string[],
   vectorStoreEnabled: boolean = false,
+  userSystemPrompt: string = "",
 ): string {
   let prompt = `You are an AI assistant in a collaborative notebook environment.
 
@@ -150,6 +151,10 @@ After executing code cells you should review the code and make changes to improv
 
 `;
 
+  if (userSystemPrompt) {
+    prompt += `\n\n${userSystemPrompt}\n`;
+  }
+
   const vectorStoreExtras =
     `IMPORTANT: If you have access to vector store tools (query_documents, find_mounted_file,
 list_indexed_files),
@@ -529,6 +534,7 @@ export async function executeAI(
   sessionId: string,
   notebookTools: NotebookTool[] = [],
   maxIterations: number = 10,
+  userSystemPrompt: string = "",
 ): Promise<{ success: boolean; error?: string }> {
   const {
     cell,
@@ -587,6 +593,7 @@ export async function executeAI(
           cell.id,
           extractedFilePaths,
           isVectorStoreIndexingEnabled(),
+          userSystemPrompt,
         ),
         prompt,
       );

packages/pyodide-runtime-agent/src/pyodide-agent.ts

@@ -22,6 +22,7 @@ import {
   isTextBasedMimeType,
   KNOWN_MIME_TYPES,
   type KnownMimeType,
+  tables,
 } from "@runt/schema";
 import {
   getBootstrapPackages as _getBootstrapPackages,
@@ -551,13 +552,26 @@ export class PyodideRuntimeAgent extends RuntimeAgent {
 
       try {
         const maxIterations = this.config.aiMaxIterations;
+
+        let userSystemPrompt = this.store.query(
+          tables.notebookMetadata
+            .select()
+            .where({ key: "user_system_prompt" })
+            .first({ fallback: () => "" }),
+        );
+
+        if (typeof userSystemPrompt !== "string") {
+          userSystemPrompt = userSystemPrompt.value;
+        }
+
         return await executeAI(
           aiContext,
           notebookContext,
           this.store,
           this.config.sessionId,
           notebookTools,
           maxIterations,
+          userSystemPrompt,
         );
       } finally {
         this.currentAIExecution = null;