chore(text_to_sql): replace user prompt renderer w/ simple text param

timescale · Feb 13, 2025 · ac3ebcf · ac3ebcf
1 parent 925f1e2
commit ac3ebcf
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 95 deletions.
diff --git a/projects/extension/sql/idempotent/906-text-to-sql-anthropic.sql b/projects/extension/sql/idempotent/906-text-to-sql-anthropic.sql
@@ -1,37 +1,5 @@
 --FEATURE-FLAG: text_to_sql
 
-
--------------------------------------------------------------------------------
--- text_to_sql_render_prompt
-create or replace function ai.text_to_sql_render_prompt
-( question text
-, obj_prompt text
-, samples_sql text
-, sql_prompt text
-) returns text
-as $func$
-    select concat_ws
-    ( E'\n'
-    , $$Below are descriptions of database objects and examples of SQL statements that are meant to give context to a user's question.$$
-    , $$Analyze the context provided. Identify the elements that are relevant to the user's question.$$
-    , $$Do not alias columns unless it is necessary.$$
-    , $$ALWAYS fully qualify references to database objects.$$
-    , $$ONLY use database elements that have been described to you. If more context is needed, use the "request_more_context_by_question" tool to ask questions about the database model.$$
-    , $$Pay attention to implicit data formats. If needed, use the "request_table_sample" to observe what actual values look like.$$
-    , $$If enough context has been provided to confidently address the question, use the "answer_user_question_with_sql_statement" tool to record your final answer in the form of a valid SQL statement.$$
-    , E'\n'
-    , coalesce(obj_prompt, '')
-    , E'\n'
-    , coalesce(samples_sql, '')
-    , E'\n'
-    , coalesce(sql_prompt, '')
-    , E'\n'
-    , concat('Q: ', question)
-    )
-$func$ language sql immutable security invoker
-set search_path to pg_catalog, pg_temp
-;
-
 -------------------------------------------------------------------------------
 -- text_to_sql_anthropic
 create or replace function ai.text_to_sql_anthropic
@@ -52,7 +20,7 @@ create or replace function ai.text_to_sql_anthropic
 , max_iter int2 default null
 , obj_renderer regprocedure default null
 , sql_renderer regprocedure default null
-, prompt_renderer regprocedure default null
+, user_prompt text default null
 , system_prompt text default null
 ) returns jsonb
 as $func$
@@ -76,15 +44,14 @@ as $func$
     , 'max_iter': max_iter
     , 'obj_renderer': obj_renderer
     , 'sql_renderer': sql_renderer
-    , 'prompt_renderer': prompt_renderer
+    , 'user_prompt': user_prompt
     , 'system_prompt': system_prompt
     absent on null
     )
 $func$ language sql immutable security invoker
 set search_path to pg_catalog, pg_temp
 ;
 
-
 -------------------------------------------------------------------------------
 -- _text_to_sql_anthropic
 create function ai._text_to_sql_anthropic
@@ -102,7 +69,6 @@ declare
     _max_vector_dist float8;
     _obj_renderer regprocedure;
     _sql_renderer regprocedure;
-    _prompt_renderer regprocedure;
     _model text;
     _max_tokens int4;
     _api_key text;
@@ -124,6 +90,7 @@ declare
     _prompt_sql text;
     _samples jsonb = '{}';
     _samples_sql text;
+    _prompt_header text;
     _prompt text;
     _system_prompt text;
     _tools jsonb;
@@ -139,7 +106,6 @@ begin
     _max_vector_dist = (_config->>'max_vector_dist')::float8;
     _obj_renderer = coalesce((_config->>'obj_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_obj(bigint, oid, oid)'::pg_catalog.regprocedure);
     _sql_renderer = coalesce((_config->>'sql_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_sql(bigint, text, text)'::pg_catalog.regprocedure);
-    _prompt_renderer = coalesce((_config->>'prompt_renderer')::pg_catalog.regprocedure, 'ai.text_to_sql_render_prompt(text, text, text, text)'::pg_catalog.regprocedure);
     _model = coalesce(_config->>'model', 'claude-3-5-sonnet-latest');
     _max_tokens = coalesce(_config operator(pg_catalog.->>) 'max_tokens', '1024')::int4;
     _api_key = _config operator(pg_catalog.->>) 'api_key';
@@ -160,6 +126,14 @@ begin
           , 'You have access to tools.'
           )
         );
+    _prompt_header = concat_ws
+        ( E'\n'
+        , $$Below are descriptions of database objects and examples of SQL statements that are meant to give context to a user's question.$$
+        , $$Analyze the context provided. Identify the elements that are relevant to the user's question.$$
+        , $$ONLY use database elements that have been described to you. If more context is needed, use the "request_more_context_by_question" tool to ask questions about the database model.$$
+        , $$If enough context has been provided to confidently address the question, use the "answer_user_question_with_sql_statement" tool to record your final answer in the form of a valid SQL statement.$$
+        , coalesce(_config operator(pg_catalog.->>) 'user_prompt', '')
+        );
 
     while _iter_remaining > 0 loop
         raise debug 'iteration: %', (_max_iter - _iter_remaining + 1);
@@ -287,17 +261,14 @@ begin
 
         -- render the user prompt
         raise debug 'rendering user prompt';
-        select format
-        ( $sql$select %I.%I($1, $2, $3, $4)$sql$
-        , n.nspname
-        , f.proname
-        )
-        into strict _sql
-        from pg_proc f
-        inner join pg_namespace n on (f.pronamespace = n.oid)
-        where f.oid = _prompt_renderer::oid
-        ;
-        execute _sql using question, _prompt_obj, _samples_sql, _prompt_sql into _prompt;
+        _prompt = concat_ws
+        ( E'\n'
+        , _prompt_header
+        , coalesce(_prompt_obj, '')
+        , coalesce(_samples_sql, '')
+        , coalesce(_prompt_sql, '')
+        , concat('Q: ', question)
+        );
         raise debug '%', _prompt;
 
         -- call llm -----------------------------------------------------------

diff --git a/projects/extension/sql/idempotent/907-text-to-sql-ollama.sql b/projects/extension/sql/idempotent/907-text-to-sql-ollama.sql
@@ -12,7 +12,7 @@ create or replace function ai.text_to_sql_ollama
 , max_vector_dist pg_catalog.float8 default null
 , obj_renderer pg_catalog.regprocedure default null
 , sql_renderer pg_catalog.regprocedure default null
-, prompt_renderer regprocedure default null
+, user_prompt text default null
 , system_prompt text default null
 ) returns pg_catalog.jsonb
 as $func$
@@ -27,7 +27,7 @@ as $func$
     , 'max_vector_dist': max_vector_dist
     , 'obj_renderer': obj_renderer
     , 'sql_renderer': sql_renderer
-    , 'prompt_renderer': prompt_renderer
+    , 'user_prompt': user_prompt
     , 'system_prompt': system_prompt
     absent on null
     )
@@ -52,7 +52,7 @@ declare
     _max_vector_dist float8;
     _obj_renderer regprocedure;
     _sql_renderer regprocedure;
-    _prompt_renderer regprocedure;
+    _prompt_header text;
     _model text;
     _host text;
     _keep_alive text;
@@ -80,7 +80,6 @@ begin
     _max_vector_dist = (_config->>'max_vector_dist')::float8;
     _obj_renderer = coalesce((_config->>'obj_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_obj(bigint, oid, oid)'::pg_catalog.regprocedure);
     _sql_renderer = coalesce((_config->>'sql_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_sql(bigint, text, text)'::pg_catalog.regprocedure);
-    _prompt_renderer = coalesce((_config->>'prompt_renderer')::pg_catalog.regprocedure, 'ai.text_to_sql_render_prompt(text, text, text, text)'::pg_catalog.regprocedure);
     _model = coalesce(_config->>'model', 'qwen2.5-coder');
     _host = config->>'host';
     _keep_alive = config->>'keep_alive';
@@ -93,7 +92,15 @@ begin
           , 'You have access to tools.'
           )
         );
-
+    _prompt_header = concat_ws
+        ( E'\n'
+        , $$Below are descriptions of database objects and examples of SQL statements that are meant to give context to a user's question.$$
+        , $$Analyze the context provided. Identify the elements that are relevant to the user's question.$$
+        , $$ONLY use database elements that have been described to you. If more context is needed, use the "request_more_context_by_question" tool to ask questions about the database model.$$
+        , $$If enough context has been provided to confidently address the question, use the "answer_user_question_with_sql_statement" tool to record your final answer in the form of a valid SQL statement.$$
+        , coalesce(_config operator(pg_catalog.->>) 'user_prompt', '')
+        );
+
     while _iter_remaining > 0 loop
         raise debug 'iteration: %', (_max_iter - _iter_remaining + 1);
         raise debug 'searching with % questions', jsonb_array_length(_questions);
@@ -213,17 +220,14 @@ begin
 
         -- render the user prompt
         raise debug 'rendering user prompt';
-        select format
-        ( $sql$select %I.%I($1, $2, $3, $4)$sql$
-        , n.nspname
-        , f.proname
-        )
-        into strict _sql
-        from pg_proc f
-        inner join pg_namespace n on (f.pronamespace = n.oid)
-        where f.oid = _prompt_renderer::oid
-        ;
-        execute _sql using question, _prompt_obj, '', _prompt_sql into _prompt;
+        _prompt = concat_ws
+        ( E'\n'
+        , _prompt_header
+        , coalesce(_prompt_obj, '')
+        --, coalesce(_samples_sql, '') TODO: implement table sampling
+        , coalesce(_prompt_sql, '')
+        , concat('Q: ', question)
+        );
         raise debug '%', _prompt;
 
         -- call llm -----------------------------------------------------------

diff --git a/projects/extension/sql/idempotent/908-text-to-sql-openai.sql b/projects/extension/sql/idempotent/908-text-to-sql-openai.sql
@@ -24,7 +24,7 @@ create or replace function ai.text_to_sql_openai
 , max_vector_dist pg_catalog.float8 default null
 , obj_renderer pg_catalog.regprocedure default null
 , sql_renderer pg_catalog.regprocedure default null
-, prompt_renderer pg_catalog.regprocedure default null
+, user_prompt text default null
 , system_prompt text default null
 ) returns pg_catalog.jsonb
 as $func$
@@ -51,7 +51,7 @@ as $func$
     , 'max_vector_dist': max_vector_dist
     , 'obj_renderer': obj_renderer
     , 'sql_renderer': sql_renderer
-    , 'prompt_renderer': prompt_renderer
+    , 'user_prompt': user_prompt
     , 'system_prompt': system_prompt
     absent on null
     )
@@ -76,7 +76,6 @@ declare
     _max_vector_dist float8;
     _obj_renderer regprocedure;
     _sql_renderer regprocedure;
-    _prompt_renderer regprocedure;
     _model text;
     _api_key text;
     _api_key_name text;
@@ -103,6 +102,7 @@ declare
     _samples_sql text;
     _prompt_obj text;
     _prompt_sql text;
+    _prompt_header text;
     _prompt text;
     _response jsonb;
     _message record;
@@ -117,7 +117,6 @@ begin
     _max_vector_dist = (_config->>'max_vector_dist')::float8;
     _obj_renderer = coalesce((_config->>'obj_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_obj(bigint, oid, oid)'::pg_catalog.regprocedure);
     _sql_renderer = coalesce((_config->>'sql_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_sql(bigint, text, text)'::pg_catalog.regprocedure);
-    _prompt_renderer = coalesce((_config->>'prompt_renderer')::pg_catalog.regprocedure, 'ai.text_to_sql_render_prompt(text, text, text, text)'::pg_catalog.regprocedure);
     _model = coalesce(_config->>'model', 'o3-mini');
     _api_key = _config operator(pg_catalog.->>) 'api_key';
     _api_key_name = _config operator(pg_catalog.->>) 'api_key_name';
@@ -141,6 +140,14 @@ begin
           , 'You have access to tools.'
           )
         );
+    _prompt_header = concat_ws
+        ( E'\n'
+        , $$Below are descriptions of database objects and examples of SQL statements that are meant to give context to a user's question.$$
+        , $$Analyze the context provided. Identify the elements that are relevant to the user's question.$$
+        , $$ONLY use database elements that have been described to you. If more context is needed, use the "request_more_context_by_question" tool to ask questions about the database model.$$
+        , $$If enough context has been provided to confidently address the question, use the "answer_user_question_with_sql_statement" tool to record your final answer in the form of a valid SQL statement.$$
+        , coalesce(_config operator(pg_catalog.->>) 'user_prompt', '')
+        );
 
     while _iter_remaining > 0 loop
         raise debug 'iteration: %', (_max_iter - _iter_remaining + 1);
@@ -269,17 +276,14 @@ begin
 
         -- render the user prompt
         raise debug 'rendering user prompt';
-        select format
-        ( $sql$select %I.%I($1, $2, $3, $4)$sql$
-        , n.nspname
-        , f.proname
-        )
-        into strict _sql
-        from pg_proc f
-        inner join pg_namespace n on (f.pronamespace = n.oid)
-        where f.oid = _prompt_renderer::oid
-        ;
-        execute _sql using question, _prompt_obj, _samples_sql, _prompt_sql into _prompt;
+        _prompt = concat_ws
+        ( E'\n'
+        , _prompt_header
+        , coalesce(_prompt_obj, '')
+        , coalesce(_samples_sql, '')
+        , coalesce(_prompt_sql, '')
+        , concat('Q: ', question)
+        );
         raise debug '%', _prompt;
 
         -- call llm -----------------------------------------------------------

diff --git a/projects/extension/sql/idempotent/909-text-to-sql-cohere.sql b/projects/extension/sql/idempotent/909-text-to-sql-cohere.sql
@@ -20,7 +20,7 @@ create or replace function ai.text_to_sql_cohere
 , max_vector_dist pg_catalog.float8 default null
 , obj_renderer pg_catalog.regprocedure default null
 , sql_renderer pg_catalog.regprocedure default null
-, prompt_renderer pg_catalog.regprocedure default null
+, user_prompt text default null
 , system_prompt text default null
 ) returns pg_catalog.jsonb
 as $func$
@@ -43,7 +43,7 @@ as $func$
     , 'max_vector_dist': max_vector_dist
     , 'obj_renderer': obj_renderer
     , 'sql_renderer': sql_renderer
-    , 'prompt_renderer': prompt_renderer
+    , 'user_prompt': user_prompt
     , 'system_prompt': system_prompt
     absent on null
     )
@@ -69,7 +69,6 @@ declare
     _max_vector_dist float8;
     _obj_renderer regprocedure;
     _sql_renderer regprocedure;
-    _prompt_renderer regprocedure;
     _model text;
     _api_key text;
     _api_key_name text;
@@ -90,6 +89,7 @@ declare
     _sql text;
     _prompt_obj text;
     _prompt_sql text;
+    _prompt_header text;
     _prompt text;
     _response jsonb;
     _tool_call record;
@@ -103,7 +103,6 @@ begin
     _max_vector_dist = (_config->>'max_vector_dist')::float8;
     _obj_renderer = coalesce((_config->>'obj_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_obj(bigint, oid, oid)'::pg_catalog.regprocedure);
     _sql_renderer = coalesce((_config->>'sql_renderer')::pg_catalog.regprocedure, 'ai.render_semantic_catalog_sql(bigint, text, text)'::pg_catalog.regprocedure);
-    _prompt_renderer = coalesce((_config->>'prompt_renderer')::pg_catalog.regprocedure, 'ai.text_to_sql_render_prompt(text, text, text, text)'::pg_catalog.regprocedure);
     _model = coalesce( _config->>'model', 'command-r7b-12-2024');
     _api_key = _config operator(pg_catalog.->>) 'api_key';
     _api_key_name = _config operator(pg_catalog.->>) 'api_key_name';
@@ -124,7 +123,15 @@ begin
           , 'You have access to tools.'
           )
         );
-
+    _prompt_header = concat_ws
+        ( E'\n'
+        , $$Below are descriptions of database objects and examples of SQL statements that are meant to give context to a user's question.$$
+        , $$Analyze the context provided. Identify the elements that are relevant to the user's question.$$
+        , $$ONLY use database elements that have been described to you. If more context is needed, use the "request_more_context_by_question" tool to ask questions about the database model.$$
+        , $$If enough context has been provided to confidently address the question, use the "answer_user_question_with_sql_statement" tool to record your final answer in the form of a valid SQL statement.$$
+        , coalesce(_config operator(pg_catalog.->>) 'user_prompt', '')
+        );
+
     while _iter_remaining > 0 loop
         raise debug 'iteration: %', (_max_iter - _iter_remaining + 1);
         raise debug 'searching with % questions', jsonb_array_length(_questions);
@@ -244,17 +251,14 @@ begin
 
         -- render the user prompt
         raise debug 'rendering user prompt';
-        select format
-        ( $sql$select %I.%I($1, $2, $3, $4)$sql$
-        , n.nspname
-        , f.proname
-        )
-        into strict _sql
-        from pg_proc f
-        inner join pg_namespace n on (f.pronamespace = n.oid)
-        where f.oid = _prompt_renderer::oid
-        ;
-        execute _sql using question, _prompt_obj, '', _prompt_sql into _prompt;
+        _prompt = concat_ws
+        ( E'\n'
+        , _prompt_header
+        , coalesce(_prompt_obj, '')
+        --, coalesce(_samples_sql, '') TODO: implement table sampling
+        , coalesce(_prompt_sql, '')
+        , concat('Q: ', question)
+        );
         raise debug '%', _prompt;
 
         -- call llm -----------------------------------------------------------