Skip to content

Commit 4f4cebb

Browse files
Light2DarkLight2Dark
authored
add datasource AI context (#6618)
## 📝 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). --> <img width="510" height="506" alt="CleanShot 2025-10-01 at 19 36 39" src="https://github.com/user-attachments/assets/0560c114-eb74-48da-a203-8c28ac9ca383" /> <img width="527" height="268" alt="CleanShot 2025-10-01 at 19 36 56" src="https://github.com/user-attachments/assets/2e66c6c4-07d8-44ee-9740-89068c0d6d62" /> result: <img width="423" height="630" alt="CleanShot 2025-10-01 at 21 39 25" src="https://github.com/user-attachments/assets/0519331b-565c-4afa-9259-caeb7cde4a18" /> todo: - [x] for internal engine, package all dataframes & internal tables - [x] formatContext needs work - [x] tests - [ ] new bug found (it tries to call tools even in `manual mode`. - [ ] bug(?), since we have the database tool, it will try to call that even if we pass the datasource schema. - [ ] improvement - it should write queries with the engine name ## 🔍 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). - [x] 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.
1 parent 114a798 commit 4f4cebb

File tree

13 files changed

+1033
-11
lines changed

13 files changed

+1033
-11
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
export {
44
BookMarkedIcon as IndexIcon,
5+
Cog as DatasourceIcon,
56
ColumnsIcon as ColumnIcon,
67
DatabaseIcon,
78
EyeIcon as ViewIcon,

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ describe("contextToXml", () => {
103103
);
104104
});
105105

106-
it("should handle complex nested data", () => {
106+
it("should handle json string data", () => {
107107
const context = {
108108
type: "complex",
109109
data: {
@@ -119,6 +119,37 @@ describe("contextToXml", () => {
119119
);
120120
});
121121

122+
it("should handle objects", () => {
123+
const context = {
124+
type: "object",
125+
data: {
126+
name: "test",
127+
config: { key: "value", nested: { prop: "test" } },
128+
},
129+
details: "Complex configuration data",
130+
};
131+
132+
const result = contextToXml(context);
133+
expect(result).toMatchInlineSnapshot(
134+
`"<object name="test" config="{"key":"value","nested":{"prop":"test"}}">Complex configuration data</object>"`,
135+
);
136+
});
137+
138+
it("should handle arrays", () => {
139+
const context = {
140+
type: "array",
141+
data: {
142+
name: "test",
143+
array: [1, 2, 3],
144+
},
145+
};
146+
147+
const result = contextToXml(context);
148+
expect(result).toMatchInlineSnapshot(
149+
`"<array name="test" array="[1,2,3]"></array>"`,
150+
);
151+
});
152+
122153
it("should handle boolean values", () => {
123154
const context = {
124155
type: "flags",

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,33 @@
11
/* Copyright 2024 Marimo. All rights reserved. */
22

3-
import { allTablesAtom } from "@/core/datasets/data-source-connections";
3+
import {
4+
allTablesAtom,
5+
dataSourceConnectionsAtom,
6+
} from "@/core/datasets/data-source-connections";
47
import { getRequestClient } from "@/core/network/requests";
58
import type { JotaiStore } from "@/core/state/jotai";
69
import { variablesAtom } from "@/core/variables/state";
710
import { CellOutputContextProvider } from "./providers/cell-output";
11+
import { DatasourceContextProvider } from "./providers/datasource";
812
import { ErrorContextProvider } from "./providers/error";
913
import { FileContextProvider } from "./providers/file";
1014
import { TableContextProvider } from "./providers/tables";
1115
import { VariableContextProvider } from "./providers/variable";
1216
import { AIContextRegistry } from "./registry";
1317

1418
export function getAIContextRegistry(store: JotaiStore) {
19+
const datasource = store.get(dataSourceConnectionsAtom);
1520
const tablesMap = store.get(allTablesAtom);
1621
const variables = store.get(variablesAtom);
1722

1823
return new AIContextRegistry()
1924
.register(new TableContextProvider(tablesMap))
2025
.register(new VariableContextProvider(variables, tablesMap))
2126
.register(new ErrorContextProvider(store))
22-
.register(new CellOutputContextProvider(store));
27+
.register(new CellOutputContextProvider(store))
28+
.register(
29+
new DatasourceContextProvider(datasource.connectionsMap, tablesMap),
30+
);
2331
}
2432

2533
export function getFileContextProvider(): FileContextProvider {

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
2+
3+
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>"`;
4+
5+
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>"`;
6+
7+
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>"`;

0 commit comments

Comments
 (0)