diff --git a/frontend/package.json b/frontend/package.json
index 3ad5157275b..cd9ec2f2169 100644
--- a/frontend/package.json
+++ b/frontend/package.json
@@ -78,6 +78,7 @@
     "@types/react-grid-layout": "^1.3.5",
     "@uidotdev/usehooks": "^2.4.1",
     "@uiw/codemirror-extensions-langs": "^4.23.0",
+    "@uiw/codemirror-extensions-mentions": "^4.23.0",
     "@uiw/react-codemirror": "^4.23.0",
     "@valtown/codemirror-codeium": "^1.1.1",
     "@xterm/addon-attach": "^0.11.0",
diff --git a/frontend/pnpm-lock.yaml b/frontend/pnpm-lock.yaml
index 3f7a0cd4fdb..355d30435ec 100644
--- a/frontend/pnpm-lock.yaml
+++ b/frontend/pnpm-lock.yaml
@@ -200,6 +200,9 @@ dependencies:
   '@uiw/codemirror-extensions-langs':
     specifier: ^4.23.0
     version: 4.23.0(@codemirror/autocomplete@6.18.0)(@codemirror/language-data@6.5.1)(@codemirror/language@6.10.2)(@codemirror/legacy-modes@6.4.0)(@codemirror/state@6.4.1)(@codemirror/view@6.32.0)(@lezer/common@1.2.1)(@lezer/highlight@1.2.1)(@lezer/javascript@1.4.14)(@lezer/lr@1.4.2)
+  '@uiw/codemirror-extensions-mentions':
+    specifier: ^4.23.0
+    version: 4.23.0(@codemirror/state@6.4.1)(@codemirror/view@6.32.0)
   '@uiw/react-codemirror':
     specifier: ^4.23.0
     version: 4.23.0(@babel/runtime@7.24.4)(@codemirror/autocomplete@6.18.0)(@codemirror/language@6.10.2)(@codemirror/lint@6.8.1)(@codemirror/search@6.5.6)(@codemirror/state@6.4.1)(@codemirror/theme-one-dark@6.1.2)(@codemirror/view@6.32.0)(codemirror@6.0.1)(react-dom@18.3.1)(react@18.3.1)
@@ -8765,6 +8768,16 @@ packages:
       - '@lezer/lr'
     dev: false
 
+  /@uiw/codemirror-extensions-mentions@4.23.0(@codemirror/state@6.4.1)(@codemirror/view@6.32.0):
+    resolution: {integrity: sha512-l+UwQcnpXvf/r8m3LLx27+5oB41Uh/6hv0JkOy1kHfNj8bZqGBMBgrnObGFfkC/oYOJHeOXvtCb0VFShTDcBRw==}
+    peerDependencies:
+      '@codemirror/state': '>=6.0.0'
+      '@codemirror/view': '>=6.0.0'
+    dependencies:
+      '@codemirror/state': 6.4.1
+      '@codemirror/view': 6.32.0
+    dev: false
+
   /@uiw/react-codemirror@4.23.0(@babel/runtime@7.24.4)(@codemirror/autocomplete@6.18.0)(@codemirror/language@6.10.2)(@codemirror/lint@6.8.1)(@codemirror/search@6.5.6)(@codemirror/state@6.4.1)(@codemirror/theme-one-dark@6.1.2)(@codemirror/view@6.32.0)(codemirror@6.0.1)(react-dom@18.3.1)(react@18.3.1):
     resolution: {integrity: sha512-MnqTXfgeLA3fsUUQjqjJgemEuNyoGALgsExVm0NQAllAAi1wfj+IoKFeK+h3XXMlTFRCFYOUh4AHDv0YXJLsOg==}
     peerDependencies:
diff --git a/frontend/src/components/editor/ai/add-cell-with-ai.tsx b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
index 3a6af39ab12..d08e9f2c5a4 100644
--- a/frontend/src/components/editor/ai/add-cell-with-ai.tsx
+++ b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
@@ -2,17 +2,83 @@
 import { useCellActions } from "../../../core/cells/cells";
 import { cn } from "@/utils/cn";
 import { Button } from "@/components/ui/button";
-import { Loader2Icon, SparklesIcon, XIcon } from "lucide-react";
+import { ChevronsUpDown, Loader2Icon, SparklesIcon, XIcon } from "lucide-react";
 import { toast } from "@/components/ui/use-toast";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import { API } from "@/core/network/api";
 import { prettyError } from "@/utils/errors";
 import { useCompletion } from "ai/react";
-import ReactCodeMirror, { EditorView } from "@uiw/react-codemirror";
+import ReactCodeMirror, {
+  EditorView,
+  keymap,
+  minimalSetup,
+} from "@uiw/react-codemirror";
+import { Prec } from "@codemirror/state";
 import { customPythonLanguageSupport } from "@/core/codemirror/language/python";
 import { asURL } from "@/utils/url";
+import { mentions } from "@uiw/codemirror-extensions-mentions";
+import { useMemo, useState } from "react";
+import { store } from "@/core/state/jotai";
+import { datasetTablesAtom } from "@/core/datasets/state";
+import { Logger } from "@/utils/Logger";
+import { Maps } from "@/utils/maps";
+import type { DataTable } from "@/core/kernel/messages";
+import { useAtom, useAtomValue } from "jotai";
+import type { Completion } from "@codemirror/autocomplete";
+import {
+  DropdownMenu,
+  DropdownMenuTrigger,
+  DropdownMenuContent,
+  DropdownMenuItem,
+} from "@/components/ui/dropdown-menu";
+import { sql } from "@codemirror/lang-sql";
+import { SQLLanguageAdapter } from "@/core/codemirror/language/sql";
+import { atomWithStorage } from "jotai/utils";
 
-const extensions = [customPythonLanguageSupport(), EditorView.lineWrapping];
+const pythonExtensions = [
+  customPythonLanguageSupport(),
+  EditorView.lineWrapping,
+];
+const sqlExtensions = [sql(), EditorView.lineWrapping];
+
+function getCompletionBody(input: string): object {
+  const datasets = extractDatasets(input);
+  Logger.debug("Included datasets", datasets);
+
+  return {
+    includeOtherCode: getCodes(""),
+    context: {
+      schema: datasets.map((dataset) => ({
+        name: dataset.name,
+        columns: dataset.columns.map((column) => ({
+          name: column.name,
+          type: column.type,
+        })),
+      })),
+    },
+    code: "",
+  };
+}
+
+function extractDatasets(input: string): DataTable[] {
+  const datasets = store.get(datasetTablesAtom);
+  const existingDatasets = Maps.keyBy(datasets, (dataset) => dataset.name);
+
+  // Extract dataset mentions from the input
+  const mentionedDatasets = input.match(/@(\w+)/g) || [];
+
+  // Filter to only include datasets that exist
+  return mentionedDatasets
+    .map((mention) => mention.slice(1))
+    .map((name) => existingDatasets.get(name))
+    .filter(Boolean);
+}
+
+// Persist across sessions
+const languageAtom = atomWithStorage<"python" | "sql">(
+  "marimo:ai-language",
+  "python",
+);
 
 /**
  * Add a cell with AI.
@@ -21,6 +87,8 @@ export const AddCellWithAI: React.FC<{
   onClose: () => void;
 }> = ({ onClose }) => {
   const { createNewCell } = useCellActions();
+  const [completionBody, setCompletionBody] = useState<object>({});
+  const [language, setLanguage] = useAtom(languageAtom);
 
   const {
     completion,
@@ -28,15 +96,15 @@ export const AddCellWithAI: React.FC<{
     stop,
     isLoading,
     setCompletion,
-    handleInputChange,
+    setInput,
     handleSubmit,
   } = useCompletion({
     api: asURL("api/ai/completion").toString(),
     headers: API.headers(),
     streamMode: "text",
     body: {
-      includeOtherCode: getCodes(""),
-      code: "",
+      ...completionBody,
+      language: language,
     },
     onError: (error) => {
       toast({
@@ -47,25 +115,40 @@ export const AddCellWithAI: React.FC<{
   });
 
   const inputComponent = (
-    <div className="flex items-center gap-3 px-3 py-2">
+    <div className="flex items-center px-3">
       <SparklesIcon className="size-4 text-[var(--blue-11)]" />
-      <input
-        className="h-8 outline-none text-base focus-visible:shadow-none flex-1 rounded-none border-none focus:border-none"
-        autoFocus={true}
+      <DropdownMenu modal={false}>
+        <DropdownMenuTrigger asChild={true}>
+          <Button
+            variant="text"
+            className="ml-2"
+            size="xs"
+            data-testid="language-button"
+          >
+            {language === "python" ? "Python" : "SQL"}
+            <ChevronsUpDown className="ml-1 h-3.5 w-3.5 text-muted-foreground/70" />
+          </Button>
+        </DropdownMenuTrigger>
+        <DropdownMenuContent align="center">
+          <DropdownMenuItem onClick={() => setLanguage("python")}>
+            Python
+          </DropdownMenuItem>
+          <DropdownMenuItem onClick={() => setLanguage("sql")}>
+            SQL
+          </DropdownMenuItem>
+        </DropdownMenuContent>
+      </DropdownMenu>
+      <PromptInput
+        onClose={() => {
+          setCompletion("");
+          onClose();
+        }}
         value={input}
-        onChange={handleInputChange}
-        onKeyDown={(e) => {
-          if (e.key === "Enter") {
-            e.preventDefault();
-            handleSubmit(e as unknown as React.FormEvent<HTMLFormElement>);
-          }
-          if (e.key === "Escape") {
-            e.preventDefault();
-            setCompletion("");
-            onClose();
-          }
+        onChange={(newValue) => {
+          setInput(newValue);
+          setCompletionBody(getCompletionBody(newValue));
         }}
-        placeholder="Generate with AI"
+        onSubmit={handleSubmit}
       />
       {isLoading && (
         <Button
@@ -90,7 +173,10 @@ export const AddCellWithAI: React.FC<{
             createNewCell({
               cellId: "__end__",
               before: false,
-              code: completion,
+              code:
+                language === "python"
+                  ? completion
+                  : SQLLanguageAdapter.fromQuery(completion),
             });
             setCompletion("");
             onClose();
@@ -106,16 +192,142 @@ export const AddCellWithAI: React.FC<{
   );
 
   return (
-    <div className={cn("flex flex-col w-full")}>
+    <div className={cn("flex flex-col w-full gap-2 py-2")}>
       {inputComponent}
+      {!completion && (
+        <span className="text-xs text-muted-foreground px-3">
+          You can mention{" "}
+          <span className="text-[var(--cyan-11)]">@dataframe</span> or{" "}
+          <span className="text-[var(--cyan-11)]">@sql_table</span> to pull
+          additional context such as column names.
+        </span>
+      )}
       {completion && (
         <ReactCodeMirror
           value={completion}
           className="cm border-t"
           onChange={setCompletion}
-          extensions={extensions}
+          extensions={language === "python" ? pythonExtensions : sqlExtensions}
         />
       )}
     </div>
   );
 };
+
+interface PromptInputProps {
+  value: string;
+  onClose: () => void;
+  onChange: (value: string) => void;
+  onSubmit: () => void;
+}
+
+const PromptInput = ({
+  value,
+  onChange,
+  onSubmit,
+  onClose,
+}: PromptInputProps) => {
+  const handleSubmit = onSubmit;
+  const handleEscape = onClose;
+  const tables = useAtomValue(datasetTablesAtom);
+
+  const extensions = useMemo(() => {
+    const completions = tables.map(
+      (table): Completion => ({
+        label: `@${table.name}`,
+        info: () => {
+          const shape = [
+            table.num_rows == null ? undefined : `${table.num_rows} rows`,
+            table.num_columns == null
+              ? undefined
+              : `${table.num_columns} columns`,
+          ]
+            .filter(Boolean)
+            .join(", ");
+
+          const infoContainer = document.createElement("div");
+          infoContainer.classList.add("prose", "prose-sm", "dark:prose-invert");
+
+          if (shape) {
+            const shapeElement = document.createElement("div");
+            shapeElement.textContent = shape;
+            shapeElement.style.fontWeight = "bold";
+            infoContainer.append(shapeElement);
+          }
+
+          if (table.source) {
+            const sourceElement = document.createElement("figcaption");
+            sourceElement.textContent = `Source: ${table.source}`;
+            infoContainer.append(sourceElement);
+          }
+
+          if (table.columns) {
+            const columnsTable = document.createElement("table");
+            const headerRow = columnsTable.insertRow();
+            const nameHeader = headerRow.insertCell();
+            nameHeader.textContent = "Column";
+            nameHeader.style.fontWeight = "bold";
+            const typeHeader = headerRow.insertCell();
+            typeHeader.textContent = "Type";
+            typeHeader.style.fontWeight = "bold";
+
+            table.columns.forEach((column) => {
+              const row = columnsTable.insertRow();
+              const nameCell = row.insertCell();
+              nameCell.textContent = column.name;
+              const typeCell = row.insertCell();
+              typeCell.textContent = column.type;
+            });
+
+            infoContainer.append(columnsTable);
+          }
+
+          return infoContainer;
+        },
+      }),
+    );
+
+    return [
+      mentions(completions),
+      EditorView.lineWrapping,
+      minimalSetup(),
+      Prec.highest(
+        keymap.of([
+          {
+            key: "Enter",
+            preventDefault: true,
+            stopPropagation: true,
+            run: () => {
+              handleSubmit();
+              return true;
+            },
+          },
+        ]),
+      ),
+      keymap.of([
+        {
+          key: "Escape",
+          preventDefault: true,
+          stopPropagation: true,
+          run: () => {
+            handleEscape();
+            return true;
+          },
+        },
+      ]),
+    ];
+  }, [tables, handleSubmit, handleEscape]);
+
+  return (
+    <ReactCodeMirror
+      className="flex-1 font-sans"
+      autoFocus={true}
+      width="100%"
+      value={value}
+      basicSetup={false}
+      extensions={extensions}
+      onChange={onChange}
+      placeholder={"Generate with AI"}
+    />
+  );
+};
diff --git a/frontend/src/components/editor/renderers/CellArray.tsx b/frontend/src/components/editor/renderers/CellArray.tsx
index 627aa887e44..47d8e3d1e9f 100644
--- a/frontend/src/components/editor/renderers/CellArray.tsx
+++ b/frontend/src/components/editor/renderers/CellArray.tsx
@@ -252,7 +252,8 @@ const AddCellButtons: React.FC = () => {
         className={cn(
           "shadow-sm border border-border rounded transition-all duration-200 overflow-hidden divide-x divide-border flex",
           !isAiButtonOpen && "w-fit",
-          isAiButtonOpen && "opacity-100 w-full max-w-4xl shadow-lg",
+          isAiButtonOpen &&
+            "opacity-100 w-full max-w-4xl shadow-lg shadow-[var(--blue-3)]",
         )}
       >
         {renderBody()}
diff --git a/frontend/src/core/codemirror/language/sql.ts b/frontend/src/core/codemirror/language/sql.ts
index 7e2e3d4b431..c851a2548c7 100644
--- a/frontend/src/core/codemirror/language/sql.ts
+++ b/frontend/src/core/codemirror/language/sql.ts
@@ -51,6 +51,7 @@ const regexes = pairs.map(
 export class SQLLanguageAdapter implements LanguageAdapter {
   readonly type = "sql";
   readonly defaultCode = `_df = mo.sql(f"""SELECT * FROM """)`;
+  static fromQuery = (query: string) => `_df = mo.sql(f"""${query.trim()}""")`;
 
   dataframeName = "_df";
   lastQuotePrefix: QuotePrefixKind = "f";
diff --git a/frontend/src/css/app/codemirror.css b/frontend/src/css/app/codemirror.css
index 2646c7031ec..239851b40b4 100644
--- a/frontend/src/css/app/codemirror.css
+++ b/frontend/src/css/app/codemirror.css
@@ -30,10 +30,9 @@
 /* -- Tooltips: code completion and type hints -- */
 
 #root .cm-tooltip {
-  border: none;
   border-radius: 4px;
 
-  @apply bg-popover shadow-smSolid shadow-shade;
+  @apply bg-popover shadow-sm shadow-shade border-border;
 }
 
 /*
diff --git a/marimo/_server/ai/prompts.py b/marimo/_server/ai/prompts.py
new file mode 100644
index 00000000000..d40f730726f
--- /dev/null
+++ b/marimo/_server/ai/prompts.py
@@ -0,0 +1,64 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+from typing import Optional
+
+from marimo._server.models.completion import AiCompletionContext, Language
+
+
+class Prompter:
+    def __init__(
+        self,
+        code: str,
+        include_other_code: str,
+        context: Optional[AiCompletionContext],
+    ):
+        self.code = code
+        self.include_other_code = include_other_code
+        self.context = context
+
+    @staticmethod
+    def get_system_prompt(language: Language) -> str:
+        language_rules = {
+            "python": [
+                "You can only output python code.",
+                "Do not describe the code, just write the code.",
+                "Do not output markdown or backticks.",
+                "When using matplotlib to show plots, use plt.gca() instead of plt.show().",  # noqa: E501
+                "If an import already exists, do not import it again.",
+                "If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.",  # noqa: E501
+            ],
+            "markdown": ["You can only output markdown."],
+            "sql": [
+                "You can only output sql.",
+                "The SQL must use duckdb syntax.",
+            ],
+        }
+
+        if language in language_rules:
+            rules = "\n".join(
+                f"{i+1}. {rule}"
+                for i, rule in enumerate(language_rules[language])
+            )
+            return (
+                f"You are a helpful assistant that can answer questions about {language}."  # noqa: E501
+                f" Here are your rules: \n{rules}"
+            )
+        else:
+            return "You are a helpful assistant that can answer questions."
+
+    def get_prompt(self, user_prompt: str) -> str:
+        prompt = user_prompt
+        if self.include_other_code:
+            prompt += f"\n\nCode from other cells:\n{self.include_other_code}"
+        if self.code.strip():
+            prompt += f"\n\nCurrent code:\n{self.code}"
+        if self.context and self.context.schema:
+            schema_info = "\n\nAvailable schema:\n"
+            for schema in self.context.schema:
+                columns = ", ".join(
+                    [f"{col.name} ({col.type})" for col in schema.columns]
+                )
+                schema_info += f"- {schema.name}: {columns}\n"
+            prompt += schema_info
+        return prompt
diff --git a/marimo/_server/api/endpoints/ai.py b/marimo/_server/api/endpoints/ai.py
index 93f8a4e39bf..89d963127ca 100644
--- a/marimo/_server/api/endpoints/ai.py
+++ b/marimo/_server/api/endpoints/ai.py
@@ -9,6 +9,7 @@
 
 from marimo import _loggers
 from marimo._config.config import MarimoConfig
+from marimo._server.ai.prompts import Prompter
 from marimo._server.api.deps import AppState
 from marimo._server.api.status import HTTPStatus
 from marimo._server.api.utils import parse_request
@@ -16,7 +17,6 @@
     AiCompletionRequest,
 )
 from marimo._server.router import APIRouter
-from marimo._utils.assert_never import assert_never
 
 if TYPE_CHECKING:
     from anthropic import (  # type: ignore[import-not-found]
@@ -225,38 +225,9 @@ async def ai_completion(
     config = app_state.config_manager.get_config(hide_secrets=False)
     body = await parse_request(request, cls=AiCompletionRequest)
 
-    if body.language == "python":
-        system_prompt = (
-            "You are a helpful assistant that can answer questions "
-            "about python code. You can only output python code. "
-            "1. Do not describe the code, just write the code."
-            "2. Do not output markdown or backticks."
-            "3. When using matplotlib to show plots,"
-            "use plt.gca() instead of plt.show()."
-            "4. If an import already exists, do not import it again."
-            "5. If a variable is already defined, use another name, or"
-            "make it private by adding an underscore at the beginning."
-        )
-    elif body.language == "markdown":
-        system_prompt = (
-            "You are a helpful assistant that can answer questions "
-            "about markdown. You can only output markdown."
-        )
-    elif body.language == "sql":
-        system_prompt = (
-            "You are a helpful assistant that can answer questions "
-            "about sql. You can only output sql."
-        )
-    else:
-        assert_never(body.language)
-
-    prompt = body.prompt
-    if body.include_other_code:
-        prompt = (
-            f"{prompt}\n\nCode from other cells:\n{body.include_other_code}"
-        )
-    if body.code.strip():
-        prompt = f"{prompt}\n\nCurrent code:\n{body.code}"
+    prompter = Prompter(body.code, body.include_other_code, body.context)
+    system_prompt = Prompter.get_system_prompt(body.language)
+    prompt = prompter.get_prompt(body.prompt)
 
     model = get_model(config)
 
diff --git a/marimo/_server/models/completion.py b/marimo/_server/models/completion.py
index cbc70b43597..816edd1ff35 100644
--- a/marimo/_server/models/completion.py
+++ b/marimo/_server/models/completion.py
@@ -1,8 +1,28 @@
 # Copyright 2024 Marimo. All rights reserved.
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Literal
+from dataclasses import dataclass, field
+from typing import List, Literal, Optional
+
+
+@dataclass
+class SchemaColumn:
+    name: str
+    type: str
+
+
+@dataclass
+class SchemaTable:
+    name: str
+    columns: List[SchemaColumn]
+
+
+@dataclass
+class AiCompletionContext:
+    schema: List[SchemaTable] = field(default_factory=list)
+
+
+Language = Literal["python", "markdown", "sql"]
 
 
 @dataclass
@@ -10,4 +30,5 @@ class AiCompletionRequest:
     prompt: str
     include_other_code: str
     code: str
-    language: Literal["python", "markdown", "sql"] = "python"
+    context: Optional[AiCompletionContext] = None
+    language: Language = "python"
diff --git a/marimo/_smoke_tests/dataframe.py b/marimo/_smoke_tests/dataframe.py
index 1030390d7c7..927cb3e02ca 100644
--- a/marimo/_smoke_tests/dataframe.py
+++ b/marimo/_smoke_tests/dataframe.py
@@ -1,8 +1,19 @@
-# Copyright 2024 Marimo. All rights reserved.
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "dask",
+#     "vega-datasets",
+#     "polars",
+#     "altair",
+#     "pyarrow",
+#     "marimo",
+#     "pandas",
+# ]
+# ///
 
 import marimo
 
-__generated_with = "0.8.0"
+__generated_with = "0.8.7"
 app = marimo.App(width="full")
 
 
@@ -296,5 +307,33 @@ def __():
     return alt, cars, mo, pa, pd, pl, vega_datasets
 
 
+app._unparsable_cell(
+    r"""
+    SELECT * FROM cars WHERE Cylinders > 6;
+    """,
+    name="__"
+)
+
+
+@app.cell
+def __(mo):
+    _df = mo.sql(
+        f"""
+        SELECT * FROM
+        """
+    )
+    return
+
+
+@app.cell
+def __(cars, mo):
+    _df = mo.sql(
+        f"""
+        SELECT * FROM cars WHERE Cylinders > 6;
+        """
+    )
+    return
+
+
 if __name__ == "__main__":
     app.run()
diff --git a/tests/_server/ai/snapshots/system_prompts.txt b/tests/_server/ai/snapshots/system_prompts.txt
new file mode 100644
index 00000000000..abc8ed8ab94
--- /dev/null
+++ b/tests/_server/ai/snapshots/system_prompts.txt
@@ -0,0 +1,26 @@
+
+
+==================== python ====================
+
+You are a helpful assistant that can answer questions about python. Here are your rules: 
+1. You can only output python code.
+2. Do not describe the code, just write the code.
+3. Do not output markdown or backticks.
+4. When using matplotlib to show plots, use plt.gca() instead of plt.show().
+5. If an import already exists, do not import it again.
+6. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
+
+==================== markdown ====================
+
+You are a helpful assistant that can answer questions about markdown. Here are your rules: 
+1. You can only output markdown.
+
+==================== sql ====================
+
+You are a helpful assistant that can answer questions about sql. Here are your rules: 
+1. You can only output sql.
+2. The SQL must use duckdb syntax.
+
+==================== idk ====================
+
+You are a helpful assistant that can answer questions.
\ No newline at end of file
diff --git a/tests/_server/ai/snapshots/user_prompts.txt b/tests/_server/ai/snapshots/user_prompts.txt
new file mode 100644
index 00000000000..33128a32e44
--- /dev/null
+++ b/tests/_server/ai/snapshots/user_prompts.txt
@@ -0,0 +1,47 @@
+
+
+==================== no code ====================
+
+Create a pandas dataframe
+
+==================== with code ====================
+
+Create a pandas dataframe
+
+Current code:
+df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
+
+==================== with code and other code ====================
+
+Create a pandas dataframe
+
+Code from other cells:
+import pandas as pd
+import numpy as np
+
+
+Current code:
+df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
+
+==================== with just other code ====================
+
+Create a pandas dataframe
+
+Code from other cells:
+import pandas as pd
+import numpy as np
+
+
+==================== with context ====================
+
+Create a pandas dataframe
+
+Code from other cells:
+import marimo as mo
+
+Current code:
+import pandas as pd
+
+Available schema:
+- df_1: age (int), name (str)
+- d2_2: a (int), b (int)
diff --git a/tests/_server/ai/test_prompts.py b/tests/_server/ai/test_prompts.py
new file mode 100644
index 00000000000..e5e062c4c75
--- /dev/null
+++ b/tests/_server/ai/test_prompts.py
@@ -0,0 +1,85 @@
+# Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
+from typing import cast
+
+from marimo._ast.visitor import Language
+from marimo._server.ai.prompts import Prompter
+from marimo._server.models.completion import (
+    AiCompletionContext,
+    SchemaColumn,
+    SchemaTable,
+)
+from tests.mocks import snapshotter
+
+snapshot = snapshotter(__file__)
+
+
+def _header(title: str) -> str:
+    return f"\n\n==================== {title} ====================\n\n"
+
+
+def test_system_prompts():
+    result = ""
+    for language in ("python", "markdown", "sql", "idk"):
+        result += _header(language)
+        result += Prompter.get_system_prompt(cast(Language, language))
+
+    snapshot("system_prompts.txt", result)
+
+
+def test_user_prompts():
+    prompt = "Create a pandas dataframe"
+
+    result = ""
+    result += _header("no code")
+    result += Prompter(
+        code="", include_other_code="", context=AiCompletionContext()
+    ).get_prompt(prompt)
+
+    result += _header("with code")
+    result += Prompter(
+        code="df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})",
+        include_other_code="",
+        context=AiCompletionContext(),
+    ).get_prompt(prompt)
+
+    result += _header("with code and other code")
+    result += Prompter(
+        code="df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})",
+        include_other_code="import pandas as pd\nimport numpy as np\n",
+        context=AiCompletionContext(),
+    ).get_prompt(prompt)
+
+    result += _header("with just other code")
+    result += Prompter(
+        code="",
+        include_other_code="import pandas as pd\nimport numpy as np\n",
+        context=AiCompletionContext(),
+    ).get_prompt(prompt)
+
+    result += _header("with context")
+    result += Prompter(
+        code="import pandas as pd",
+        include_other_code="import marimo as mo",
+        context=AiCompletionContext(
+            schema=[
+                SchemaTable(
+                    name="df_1",
+                    columns=[
+                        SchemaColumn("age", "int"),
+                        SchemaColumn("name", "str"),
+                    ],
+                ),
+                SchemaTable(
+                    name="d2_2",
+                    columns=[
+                        SchemaColumn("a", "int"),
+                        SchemaColumn("b", "int"),
+                    ],
+                ),
+            ],
+        ),
+    ).get_prompt(prompt)
+
+    snapshot("user_prompts.txt", result)