1+ import { ActionDefinition , ActionContext , OutputObject } from 'connery' ;
2+ import pkg from 'pg' ;
3+ const { Client } = pkg ;
4+ import { Anthropic } from '@anthropic-ai/sdk' ;
5+
6+ const actionDefinition : ActionDefinition = {
7+ key : 'chatWithYourDb' ,
8+ name : 'Chat with your DB' ,
9+ description : 'Users can send DB requests in natural language and receive data and/or helpful feedback.' ,
10+ type : 'read' ,
11+ inputParameters : [
12+ {
13+ key : 'anthropicApiKey' ,
14+ name : 'Anthropic API Key' ,
15+ description : 'Your Anthropic API key' ,
16+ type : 'string' ,
17+ validation : {
18+ required : true ,
19+ } ,
20+ } ,
21+ {
22+ key : 'connectionString' ,
23+ name : 'Database Connection String' ,
24+ description : 'PostgreSQL connection string (should use read-only credentials)' ,
25+ type : 'string' ,
26+ validation : {
27+ required : true ,
28+ } ,
29+ } ,
30+ {
31+ key : 'instructions' ,
32+ name : 'Instructions' ,
33+ description : 'Optional instructions for processing the response' ,
34+ type : 'string' ,
35+ validation : {
36+ required : false ,
37+ } ,
38+ } ,
39+ {
40+ key : 'maxRows' ,
41+ name : 'Maximum Rows' ,
42+ description : 'Maximum number of rows to return (default: 100)' ,
43+ type : 'string' ,
44+ validation : {
45+ required : false ,
46+ } ,
47+ } ,
48+ {
49+ key : 'question' ,
50+ name : 'Question' ,
51+ description : 'Your database question in natural language' ,
52+ type : 'string' ,
53+ validation : {
54+ required : true ,
55+ } ,
56+ } ,
57+ ] ,
58+ operation : {
59+ handler : handler ,
60+ } ,
61+ outputParameters : [
62+ {
63+ key : 'data' ,
64+ name : 'Data' ,
65+ description : 'The data returned by your database query' ,
66+ type : 'string' ,
67+ validation : {
68+ required : true ,
69+ } ,
70+ } ,
71+ {
72+ key : 'query' ,
73+ name : 'Query' ,
74+ description : 'The generated SQL query' ,
75+ type : 'string' ,
76+ validation : {
77+ required : true ,
78+ } ,
79+ } ,
80+ ] ,
81+ } ;
82+
83+ export default actionDefinition ;
84+
85+ export async function handler ( { input } : ActionContext ) : Promise < OutputObject > {
86+ let client : pkg . Client | null = null ;
87+
88+ try {
89+ // Always generate new schema
90+ client = new Client ( input . connectionString ) ;
91+ await client . connect ( ) ;
92+ await client . query ( 'SELECT 1' ) ; // Test connection
93+ const schemaInfo = await getSchemaInfo ( client ) ;
94+
95+ const sqlQuery = await generateSqlQuery ( input . anthropicApiKey , schemaInfo , input . question , parseInt ( input . maxRows || '100' ) ) ;
96+ const result = await client . query ( sqlQuery ) ;
97+
98+ // Format each part separately
99+ const dataResponse = formatDataResponse ( result . rows , input . instructions ) ;
100+ const queryResponse = formatQueryResponse ( sqlQuery ) ;
101+
102+ // Return all responses
103+ return {
104+ data : dataResponse ,
105+ query : queryResponse ,
106+ } ;
107+ } catch ( error : unknown ) {
108+ throw error ;
109+ } finally {
110+ if ( client ) {
111+ try {
112+ await client . end ( ) ;
113+ } catch ( closeError ) {
114+ // Silently handle connection closing errors
115+ }
116+ }
117+ }
118+ }
119+
120+ async function getSchemaInfo ( client : pkg . Client ) : Promise < string > {
121+ const schemaQuery = `
122+ WITH columns_info AS (
123+ SELECT
124+ c.table_schema,
125+ c.table_name,
126+ c.column_name,
127+ c.data_type,
128+ c.is_nullable,
129+ c.column_default,
130+ c.ordinal_position
131+ FROM
132+ information_schema.columns c
133+ WHERE
134+ c.table_schema NOT IN ('pg_catalog', 'information_schema')
135+ ),
136+ primary_keys AS (
137+ SELECT
138+ kcu.table_schema,
139+ kcu.table_name,
140+ kcu.column_name
141+ FROM
142+ information_schema.table_constraints tc
143+ JOIN
144+ information_schema.key_column_usage kcu
145+ ON tc.constraint_name = kcu.constraint_name
146+ AND tc.table_schema = kcu.table_schema
147+ WHERE
148+ tc.constraint_type = 'PRIMARY KEY'
149+ AND tc.table_schema NOT IN ('pg_catalog', 'information_schema')
150+ ),
151+ foreign_keys AS (
152+ SELECT
153+ kcu.table_schema AS table_schema,
154+ kcu.table_name AS table_name,
155+ kcu.column_name AS column_name,
156+ ccu.table_schema AS foreign_table_schema,
157+ ccu.table_name AS foreign_table_name,
158+ ccu.column_name AS foreign_column_name
159+ FROM
160+ information_schema.table_constraints tc
161+ JOIN
162+ information_schema.key_column_usage kcu
163+ ON tc.constraint_name = kcu.constraint_name
164+ AND tc.table_schema = kcu.table_schema
165+ JOIN
166+ information_schema.constraint_column_usage ccu
167+ ON ccu.constraint_name = tc.constraint_name
168+ AND ccu.table_schema = tc.table_schema
169+ WHERE
170+ tc.constraint_type = 'FOREIGN KEY'
171+ AND tc.table_schema NOT IN ('pg_catalog', 'information_schema')
172+ )
173+ SELECT
174+ jsonb_pretty(
175+ jsonb_agg(
176+ jsonb_build_object(
177+ 'table_schema', tbl.table_schema,
178+ 'table_name', tbl.table_name,
179+ 'columns', tbl.columns
180+ )
181+ )
182+ ) AS schema_json
183+ FROM (
184+ SELECT
185+ c.table_schema,
186+ c.table_name,
187+ jsonb_agg(
188+ jsonb_build_object(
189+ 'column_name', c.column_name,
190+ 'data_type', c.data_type,
191+ 'is_nullable', c.is_nullable,
192+ 'column_default', c.column_default,
193+ 'is_primary_key', CASE WHEN pk.column_name IS NOT NULL THEN true ELSE false END,
194+ 'is_foreign_key', CASE WHEN fk.column_name IS NOT NULL THEN true ELSE false END,
195+ 'foreign_table_schema', fk.foreign_table_schema,
196+ 'foreign_table_name', fk.foreign_table_name,
197+ 'foreign_column_name', fk.foreign_column_name
198+ ) ORDER BY c.ordinal_position
199+ ) AS columns
200+ FROM
201+ columns_info c
202+ LEFT JOIN
203+ primary_keys pk
204+ ON c.table_schema = pk.table_schema
205+ AND c.table_name = pk.table_name
206+ AND c.column_name = pk.column_name
207+ LEFT JOIN
208+ foreign_keys fk
209+ ON c.table_schema = fk.table_schema
210+ AND c.table_name = fk.table_name
211+ AND c.column_name = fk.column_name
212+ GROUP BY
213+ c.table_schema,
214+ c.table_name
215+ ORDER BY
216+ c.table_schema,
217+ c.table_name
218+ ) tbl;
219+ ` ;
220+
221+ const schemaResult = await client . query ( schemaQuery ) ;
222+
223+ const schemaJson = schemaResult . rows [ 0 ] . schema_json ;
224+ return schemaJson ;
225+ }
226+
227+ async function generateSqlQuery ( apiKey : string , schemaInfo : string , question : string , maxRows : number ) : Promise < string > {
228+ const systemPrompt = `You are a PostgreSQL expert. Generate secure, read-only SQL queries based on natural language questions.
229+ Schema information: ${ schemaInfo }
230+
231+ Important: Return ONLY the raw SQL query without any formatting, markdown, or code blocks.
232+
233+ Rules:
234+ - Use ONLY tables and columns that exist in the provided schema information
235+ - Do not make assumptions about columns that aren't explicitly listed in the schema
236+ - Generate only SELECT queries (no INSERT, UPDATE, DELETE, etc.)
237+ - Ensure queries are optimized for performance
238+ - Include relevant JOINs when needed
239+ - Add inline comments with -- to explain the query
240+ - Limit results to ${ maxRows }
241+ - Use explicit column names instead of SELECT *
242+ - Add ORDER BY clauses when relevant
243+ - Do not include markdown code blocks or SQL syntax highlighting in your response
244+ - Do not include any other text in your response
245+ - If you cannot construct a query using only the available columns, respond with an error message starting with "ERROR:"` ;
246+
247+
248+ const ai = new Anthropic ( { apiKey } ) ;
249+ const completion = await ai . messages . create ( {
250+ model : "claude-3-5-sonnet-20241022" ,
251+ max_tokens : 8192 ,
252+ messages : [
253+ {
254+ role : "user" ,
255+ content : systemPrompt + "\n\n" + question
256+ }
257+ ] ,
258+ temperature : 0
259+ } ) ;
260+
261+ const sqlQuery = completion . content [ 0 ] ?. type === 'text' ? completion . content [ 0 ] . text : null ;
262+ if ( ! sqlQuery ) {
263+ throw new Error ( 'Failed to generate SQL query: No response from Anthropic' ) ;
264+ }
265+
266+ if ( sqlQuery . startsWith ( 'ERROR:' ) ) {
267+ throw new Error ( sqlQuery ) ;
268+ }
269+
270+ return sqlQuery ;
271+ }
272+
273+ function formatDataResponse ( rows : any [ ] , instructions ?: string ) : string {
274+ let response = '' ;
275+
276+ // Handle empty results
277+ if ( ! rows || rows . length === 0 ) {
278+ response = "No data found for your query." ;
279+ } else {
280+ try {
281+ const sanitizedRows = rows . map ( row => {
282+ const sanitizedRow : any = { } ;
283+ for ( const [ key , value ] of Object . entries ( row ) ) {
284+ sanitizedRow [ key ] = typeof value === 'bigint' || typeof value === 'number'
285+ ? value . toString ( )
286+ : value ;
287+ }
288+ return sanitizedRow ;
289+ } ) ;
290+
291+ response = JSON . stringify ( sanitizedRows , null , 2 ) ;
292+ } catch ( error ) {
293+ throw new Error ( `Error formatting database response: ${ error instanceof Error ? error . message : String ( error ) } ) ;
294+ }
295+ }
296+
297+ // Add instructions if provided
298+ if ( instructions ) {
299+ response = `Instructions for the following content: ${ instructions } ${ response } ;
300+ }
301+
302+ return response ;
303+ }
304+
305+ function formatQueryResponse ( sqlQuery : string ) : string {
306+ return sqlQuery ;
307+ }
0 commit comments