From 5e8e3cbea3622c884e6d4230a96ea284f90e6ce9 Mon Sep 17 00:00:00 2001
From: Rodrigo Rodriguez <me@rodrigorodriguez.com>
Date: Thu, 28 Nov 2024 11:13:01 -0300
Subject: [PATCH] fix(llm.gblib): Talk to data local db use fix.

---
 packages/core.gbapp/services/GBDeployer.ts        |  3 ++-
 packages/llm.gblib/services/ChatServices.ts       | 15 +++++++++------
 .../talk-to-data.gbot/config.csv                  |  3 ++-
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/packages/core.gbapp/services/GBDeployer.ts b/packages/core.gbapp/services/GBDeployer.ts
index 10558bb9..c955d679 100644
--- a/packages/core.gbapp/services/GBDeployer.ts
+++ b/packages/core.gbapp/services/GBDeployer.ts
@@ -676,8 +676,9 @@ export class GBDeployer implements IGBDeployer {
             let con = {};
             con['name'] = connectionName;
             con['storageDriver'] = min.core.getParam<string>(min.instance, `${connectionName} Driver`, null);
+            con['storageTables'] = min.core.getParam<string>(min.instance, `${connectionName} Tables`, null);
             const storageName = min.core.getParam<string>(min.instance, `${connectionName} Name`, null);
-
+            
             let file = min.core.getParam<string>(min.instance, `${connectionName} File`, null);
 
             if (storageName) {
diff --git a/packages/llm.gblib/services/ChatServices.ts b/packages/llm.gblib/services/ChatServices.ts
index 9797471e..32baf02d 100644
--- a/packages/llm.gblib/services/ChatServices.ts
+++ b/packages/llm.gblib/services/ChatServices.ts
@@ -539,6 +539,8 @@ export class ChatServices {
     } else if (LLMMode === 'sql' || LLMMode === 'chart') {
       const con = min[`llm`]['gbconnection'];
       const dialect = con['storageDriver'];
+      let tables = con['storageTables'];
+      tables = tables ? tables.split(';') : null;
 
       const answerSource = await (min.core as any)['getParam'](min.instance,
         'Answer Source', 'server');
@@ -613,14 +615,14 @@ export class ChatServices {
           SQL QUERY:`);
 
       /**
-       * Create a new RunnableSequence where we pipe the output from `db.getTableInfo()`
+       * Create a new RunnableSequence.
        */
       const sqlQueryChain = RunnableSequence.from([
         {
-          schema: async () => db.getTableInfo(),
+          schema: async () => db.getTableInfo(tables),
           question: (input: { question: string }) => input.question,
           top_k: () => 10,
-          table_info: () => 'any'
+          table_info: () => 'any',
         },
         prompt,
         model,
@@ -633,7 +635,7 @@ export class ChatServices {
        */
       const finalResponsePrompt =
         PromptTemplate.fromTemplate(`Based on the table schema below, question, SQL query, and SQL response, write a natural language response:
-          Optimize answers for KPI people. ${systemPrompt}
+          Optimize answers for KPI people. ${systemPrompt} 
               ------------
               SCHEMA: {schema}
               ------------
@@ -659,12 +661,13 @@ export class ChatServices {
           query: sqlQueryChain
         },
         {
-          schema: async () => db.getTableInfo(),
+          schema: async () => db.getTableInfo(tables),
           question: input => input.question,
           query: input => input.query,
           response: input => db.run(input.query),
           top_k: () => 10,
-          table_info: () => 'any'
+          table_info: () => 'any',
+          table_names_to_use: () => tables
         },
         {
           result: finalResponsePrompt.pipe(model).pipe(
diff --git a/templates/talk-to-data.gbai/talk-to-data.gbot/config.csv b/templates/talk-to-data.gbai/talk-to-data.gbot/config.csv
index 14b550bd..05da4925 100644
--- a/templates/talk-to-data.gbai/talk-to-data.gbot/config.csv
+++ b/templates/talk-to-data.gbai/talk-to-data.gbot/config.csv
@@ -3,4 +3,5 @@ Admin, 5521999998888
 Answer Mode,chart
 llm File,northwind.db
 llm Driver,sqlite
-Theme Color,red
\ No newline at end of file
+llm Tables,Products;Orders;Customers;Categories;Order Details
+Theme Color,red