From 84904839e59de3ac27919ff7c413d7c17ab6d47a Mon Sep 17 00:00:00 2001
From: Suman Michael <16167186+sumanmichael@users.noreply.github.com>
Date: Tue, 12 Mar 2024 07:36:33 +0530
Subject: [PATCH] added table comment to Table Info in SQLDatabase (#11774)
* added table comment in sql query
* handle SQL dialects without comment support
* typo in comments
---
.../llama_index/core/utilities/sql_wrapper.py | 17 +++++++++++++----
1 file changed, 13 insertions(+), 4 deletions(-)
diff --git a/llama-index-core/llama_index/core/utilities/sql_wrapper.py b/llama-index-core/llama_index/core/utilities/sql_wrapper.py
index b8c8a5410..dbe74bdc0 100644
--- a/llama-index-core/llama_index/core/utilities/sql_wrapper.py
+++ b/llama-index-core/llama_index/core/utilities/sql_wrapper.py
@@ -150,10 +150,19 @@ class SQLDatabase:
def get_single_table_info(self, table_name: str) -> str:
"""Get table info for a single table."""
# same logic as table_info, but with specific table names
- template = (
- "Table '{table_name}' has columns: {columns}, "
- "and foreign keys: {foreign_keys}."
- )
+ template = "Table '{table_name}' has columns: {columns}, "
+ try:
+ # try to retrieve table comment
+ table_comment = self._inspector.get_table_comment(
+ table_name, schema=self._schema
+ )["text"]
+ if table_comment:
+ template += f"with comment: ({table_comment}) "
+ except NotImplementedError:
+ # get_table_comment raises NotImplementedError for a dialect that does not support comments.
+ pass
+
+ template += "and foreign keys: {foreign_keys}."
columns = []
for column in self._inspector.get_columns(table_name, schema=self._schema):
if column.get("comment"):
--
GitLab