update to multi-cell prompt (#6640)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->

Only added for multi-cell (chat-panel, generate with AI) AI generation

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

Author: Light2Dark

marimo/_server/ai/prompts.py

@@ -16,7 +16,8 @@
 FIM_SUFFIX_TAG = "<|fim_suffix|>"
 FIM_MIDDLE_TAG = "<|fim_middle|>"
 
-language_rules = {
+LANGUAGES: list[Language] = ["python", "sql", "markdown"]
+language_rules: dict[Language, list[str]] = {
     "python": [
         "For matplotlib: use plt.gca() as the last expression instead of plt.show().",
         "For plotly: return the figure object directly.",
@@ -33,6 +34,15 @@
 }
 
 
+language_rules_multiple_cells: dict[Language, list[str]] = {
+    "sql": [
+        'SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.',
+        "This will automatically display the result in the UI. You do not need to return the dataframe in the cell.",
+        "The SQL must use the syntax of the database engine specified in the `engine` variable. If no engine, then use duckdb syntax.",
+    ]
+}
+
+
 def _format_schema_info(tables: Optional[list[SchemaTable]]) -> str:
     """Helper to format schema information from context"""
     if not tables:
@@ -166,10 +176,13 @@ def get_refactor_or_insert_notebook_cell_system_prompt(
 
     if support_multiple_cells:
         # Add all language rules for multi-cell scenarios
-        for lang in language_rules:
-            if len(language_rules[lang]) > 0:
+        for lang in LANGUAGES:
+            language_rule = language_rules_multiple_cells.get(
+                lang, language_rules.get(lang, [])
+            )
+            if language_rule:
                 system_prompt += (
-                    f"\n\n## Rules for {lang}:\n{_rules(language_rules[lang])}"
+                    f"\n\n## Rules for {lang}:\n{_rules(language_rule)}"
                 )
     elif language in language_rules and language_rules[language]:
         system_prompt += (

tests/_server/ai/snapshots/system_prompts.txt

@@ -347,7 +347,9 @@ Separate logic into multiple cells to keep the code organized and readable.
 7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
 
 ## Rules for sql:
-1. The SQL must use duckdb syntax.
+1. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
+2. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
+3. The SQL must use the syntax of the database engine specified in the `engine` variable. If no engine, then use duckdb syntax.
 
 ## Available variables from other cells:
 - variable: `df`