add datasource AI context

frontend/src/components/databases/namespace-icons.ts

@@ -2,6 +2,7 @@
 
 export {
   BookMarkedIcon as IndexIcon,
+  Cog as DatasourceIcon,
   ColumnsIcon as ColumnIcon,
   DatabaseIcon,
   EyeIcon as ViewIcon,

frontend/src/core/ai/context/__tests__/utils.test.ts

@@ -103,7 +103,7 @@ describe("contextToXml", () => {
     );
   });
 
-  it("should handle complex nested data", () => {
+  it("should handle json string data", () => {
     const context = {
       type: "complex",
       data: {
@@ -119,6 +119,37 @@ describe("contextToXml", () => {
     );
   });
 
+  it("should handle objects", () => {
+    const context = {
+      type: "object",
+      data: {
+        name: "test",
+        config: { key: "value", nested: { prop: "test" } },
+      },
+      details: "Complex configuration data",
+    };
+
+    const result = contextToXml(context);
+    expect(result).toMatchInlineSnapshot(
+      `"<object name="test" config="{"key":"value","nested":{"prop":"test"}}">Complex configuration data</object>"`,
+    );
+  });
+
+  it("should handle arrays", () => {
+    const context = {
+      type: "array",
+      data: {
+        name: "test",
+        array: [1, 2, 3],
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toMatchInlineSnapshot(
+      `"<array name="test" array="[1,2,3]"></array>"`,
+    );
+  });
+
   it("should handle boolean values", () => {
     const context = {
       type: "flags",

frontend/src/core/ai/context/context.ts

@@ -1,25 +1,33 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { allTablesAtom } from "@/core/datasets/data-source-connections";
+import {
+  allTablesAtom,
+  dataSourceConnectionsAtom,
+} from "@/core/datasets/data-source-connections";
 import { getRequestClient } from "@/core/network/requests";
 import type { JotaiStore } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import { CellOutputContextProvider } from "./providers/cell-output";
+import { DatasourceContextProvider } from "./providers/datasource";
 import { ErrorContextProvider } from "./providers/error";
 import { FileContextProvider } from "./providers/file";
 import { TableContextProvider } from "./providers/tables";
 import { VariableContextProvider } from "./providers/variable";
 import { AIContextRegistry } from "./registry";
 
 export function getAIContextRegistry(store: JotaiStore) {
+  const datasource = store.get(dataSourceConnectionsAtom);
   const tablesMap = store.get(allTablesAtom);
   const variables = store.get(variablesAtom);
 
   return new AIContextRegistry()
     .register(new TableContextProvider(tablesMap))
     .register(new VariableContextProvider(variables, tablesMap))
     .register(new ErrorContextProvider(store))
-    .register(new CellOutputContextProvider(store));
+    .register(new CellOutputContextProvider(store))
+    .register(
+      new DatasourceContextProvider(datasource.connectionsMap, tablesMap),
+    );
 }
 
 export function getFileContextProvider(): FileContextProvider {

frontend/src/core/ai/context/providers/__tests__/__snapshots__/datasource.test.ts.snap

@@ -0,0 +1,7 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`DatasourceContextProvider > formatContext > should format context for datasource with multiple databases > multi-db-context 1`] = `"<datasource connection="{"dialect":"duckdb","databases":[{"name":"db1","dialect":"duckdb","schemas":[{"name":"schema1","tables":[{"name":"table1","source_type":"connection","source":"multi","num_rows":10,"num_columns":2,"variable_name":null,"columns":[]},{"name":"table2","source_type":"connection","source":"multi","num_rows":20,"num_columns":3,"variable_name":null,"columns":[]}]}]},{"name":"db2","dialect":"duckdb","schemas":[{"name":"schema2","tables":[{"name":"table3","source_type":"connection","source":"multi","num_rows":15,"num_columns":4,"variable_name":null,"columns":[]}]},{"name":"schema3","tables":[]}]}],"default_database":"db1","default_schema":"schema1","engine_name":"multi"}"></datasource>"`;
+
+exports[`DatasourceContextProvider > formatContext > should format context for internal duckdb with multiple tables > internal-datasource-context 1`] = `"<datasource connection="{"dialect":"duckdb","databases":[{"name":"main","dialect":"duckdb","schemas":[{"name":"public","tables":[{"name":"users","source_type":"connection","source":"__marimo_duckdb","num_rows":100,"num_columns":3,"variable_name":null,"columns":[]},{"name":"orders","source_type":"connection","source":"__marimo_duckdb","num_rows":50,"num_columns":4,"variable_name":null,"columns":[]}]},{"name":"analytics","tables":[{"name":"events","source_type":"connection","source":"__marimo_duckdb","num_rows":200,"num_columns":5,"variable_name":null,"columns":[]}]}]}]}" tables="[{"name":"users","source":"local","source_type":"local","num_rows":100,"num_columns":3,"variable_name":"users","columns":[{"name":"id","type":"integer","external_type":"INTEGER","sample_values":[1,2,3]},{"name":"name","type":"string","external_type":"VARCHAR","sample_values":["Alice","Bob","Charlie"]},{"name":"age","type":"integer","external_type":"INTEGER","sample_values":[25,30,35]}],"engine":null,"indexes":null,"primary_keys":null,"type":"table"},{"name":"orders","source":"local","source_type":"local","num_rows":100,"num_columns":3,"variable_name":"orders","columns":[{"name":"id","type":"integer","external_type":"INTEGER","sample_values":[1,2,3]},{"name":"name","type":"string","external_type":"VARCHAR","sample_values":["Alice","Bob","Charlie"]},{"name":"age","type":"integer","external_type":"INTEGER","sample_values":[25,30,35]}],"engine":null,"indexes":null,"primary_keys":null,"type":"table"},{"name":"events","source":"local","source_type":"local","num_rows":100,"num_columns":3,"variable_name":"events","columns":[{"name":"id","type":"integer","external_type":"INTEGER","sample_values":[1,2,3]},{"name":"name","type":"string","external_type":"VARCHAR","sample_values":["Alice","Bob","Charlie"]},{"name":"age","type":"integer","external_type":"INTEGER","sample_values":[25,30,35]}],"engine":null,"indexes":null,"primary_keys":null,"type":"table"}]"></datasource>"`;
+
+exports[`DatasourceContextProvider > formatContext > should format context for postgres datasource > postgres-datasource-context 1`] = `"<datasource connection="{"dialect":"postgresql","databases":[{"name":"production","dialect":"postgresql","schemas":[{"name":"public","tables":[{"name":"customers","source_type":"connection","source":"postgres","num_rows":1000,"num_columns":8,"variable_name":null,"columns":[]},{"name":"products","source_type":"connection","source":"postgres","num_rows":500,"num_columns":6,"variable_name":null,"columns":[]},{"name":"sales","source_type":"connection","source":"postgres","num_rows":5000,"num_columns":10,"variable_name":null,"columns":[]}]}]}],"engine_name":"postgres"}"></datasource>"`;