QuerySUTRA 0.5.3__tar.gz → 0.6.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. {querysutra-0.5.3 → querysutra-0.6.1}/PKG-INFO +18 -2
  2. {querysutra-0.5.3 → querysutra-0.6.1}/QuerySUTRA.egg-info/PKG-INFO +18 -2
  3. querysutra-0.6.1/QuerySUTRA.egg-info/top_level.txt +1 -0
  4. querysutra-0.6.1/pyproject.toml +57 -0
  5. {querysutra-0.5.3 → querysutra-0.6.1}/requirements.txt +1 -1
  6. querysutra-0.6.1/setup.py +3 -0
  7. querysutra-0.6.1/sutra/__init__.py +6 -0
  8. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/database_manager.py +235 -195
  9. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/nlp_processor.py +175 -143
  10. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/schema_generator.py +56 -52
  11. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/sutra.py +133 -16
  12. querysutra-0.5.3/QuerySUTRA.egg-info/top_level.txt +0 -3
  13. querysutra-0.5.3/pyproject.toml +0 -17
  14. querysutra-0.5.3/setup.py +0 -3
  15. querysutra-0.5.3/sutra/__init__.py +0 -4
  16. {querysutra-0.5.3 → querysutra-0.6.1}/LICENSE +0 -0
  17. {querysutra-0.5.3 → querysutra-0.6.1}/MANIFEST.in +0 -0
  18. {querysutra-0.5.3 → querysutra-0.6.1}/QuerySUTRA.egg-info/SOURCES.txt +0 -0
  19. {querysutra-0.5.3 → querysutra-0.6.1}/QuerySUTRA.egg-info/dependency_links.txt +0 -0
  20. {querysutra-0.5.3 → querysutra-0.6.1}/QuerySUTRA.egg-info/requires.txt +0 -0
  21. {querysutra-0.5.3 → querysutra-0.6.1}/README.md +0 -0
  22. {querysutra-0.5.3 → querysutra-0.6.1}/examples/quickstart.py +0 -0
  23. {querysutra-0.5.3 → querysutra-0.6.1}/examples/sutra_usage_guide.ipynb +0 -0
  24. {querysutra-0.5.3 → querysutra-0.6.1}/examples/usage_guide.ipynb +0 -0
  25. {querysutra-0.5.3 → querysutra-0.6.1}/setup.cfg +0 -0
  26. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/cache_manager.py +0 -0
  27. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/clear_cache.py +0 -0
  28. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/core.py +0 -0
  29. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/data_loader.py +0 -0
  30. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/direct_query.py +0 -0
  31. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/feedback.py +0 -0
  32. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/feedback_matcher.py +0 -0
  33. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/schema_embeddings.py +0 -0
  34. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/sutra_client.py +0 -0
  35. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/sutra_core.py +0 -0
  36. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/sutra_simple.py +0 -0
  37. {querysutra-0.5.3 → querysutra-0.6.1}/sutra/visualizer.py +0 -0
  38. {querysutra-0.5.3 → querysutra-0.6.1}/tests/__init__.py +0 -0
  39. {querysutra-0.5.3 → querysutra-0.6.1}/tests/test_modules.py +0 -0
  40. {querysutra-0.5.3 → querysutra-0.6.1}/tests/test_sutra.py +0 -0
  41. {querysutra-0.5.3 → querysutra-0.6.1}/utils/__init__.py +0 -0
  42. {querysutra-0.5.3 → querysutra-0.6.1}/utils/file_utils.py +0 -0
  43. {querysutra-0.5.3 → querysutra-0.6.1}/utils/text_utils.py +0 -0
@@ -1,9 +1,25 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: QuerySUTRA
3
- Version: 0.5.3
4
- Summary: SUTRA
3
+ Version: 0.6.1
4
+ Summary: AI-powered data analysis for structured and unstructured data. Query PDF, Word, CSV, Excel with natural language.
5
5
  Author: Aditya Batta
6
6
  License: MIT
7
+ Project-URL: Homepage, https://github.com/adityabatta/QuerySUTRA
8
+ Project-URL: Repository, https://github.com/adityabatta/QuerySUTRA
9
+ Project-URL: Issues, https://github.com/adityabatta/QuerySUTRA/issues
10
+ Keywords: ai,data-analysis,nlp,sql,pdf,openai,natural-language,query,database
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.8
17
+ Classifier: Programming Language :: Python :: 3.9
18
+ Classifier: Programming Language :: Python :: 3.10
19
+ Classifier: Programming Language :: Python :: 3.11
20
+ Classifier: Programming Language :: Python :: 3.12
21
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
22
+ Classifier: Topic :: Database
7
23
  Requires-Python: >=3.8
8
24
  Description-Content-Type: text/markdown
9
25
  License-File: LICENSE
@@ -1,9 +1,25 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: QuerySUTRA
3
- Version: 0.5.3
4
- Summary: SUTRA
3
+ Version: 0.6.1
4
+ Summary: AI-powered data analysis for structured and unstructured data. Query PDF, Word, CSV, Excel with natural language.
5
5
  Author: Aditya Batta
6
6
  License: MIT
7
+ Project-URL: Homepage, https://github.com/adityabatta/QuerySUTRA
8
+ Project-URL: Repository, https://github.com/adityabatta/QuerySUTRA
9
+ Project-URL: Issues, https://github.com/adityabatta/QuerySUTRA/issues
10
+ Keywords: ai,data-analysis,nlp,sql,pdf,openai,natural-language,query,database
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.8
17
+ Classifier: Programming Language :: Python :: 3.9
18
+ Classifier: Programming Language :: Python :: 3.10
19
+ Classifier: Programming Language :: Python :: 3.11
20
+ Classifier: Programming Language :: Python :: 3.12
21
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
22
+ Classifier: Topic :: Database
7
23
  Requires-Python: >=3.8
8
24
  Description-Content-Type: text/markdown
9
25
  License-File: LICENSE
@@ -0,0 +1 @@
1
+ sutra
@@ -0,0 +1,57 @@
1
+ [build-system]
2
+ requires = ["setuptools>=45", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "QuerySUTRA"
7
+ version = "0.6.1"
8
+ description = "AI-powered data analysis for structured and unstructured data. Query PDF, Word, CSV, Excel with natural language."
9
+ readme = "README.md"
10
+ requires-python = ">=3.8"
11
+ license = {text = "MIT"}
12
+ authors = [{name = "Aditya Batta"}]
13
+ keywords = ["ai", "data-analysis", "nlp", "sql", "pdf", "openai", "natural-language", "query", "database"]
14
+ classifiers = [
15
+ "Development Status :: 4 - Beta",
16
+ "Intended Audience :: Developers",
17
+ "Intended Audience :: Science/Research",
18
+ "License :: OSI Approved :: MIT License",
19
+ "Programming Language :: Python :: 3",
20
+ "Programming Language :: Python :: 3.8",
21
+ "Programming Language :: Python :: 3.9",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
26
+ "Topic :: Database",
27
+ ]
28
+ dependencies = [
29
+ "pandas>=1.3.0",
30
+ "numpy>=1.21.0",
31
+ "openai>=1.0.0",
32
+ "plotly>=5.0.0",
33
+ "matplotlib>=3.3.0",
34
+ "PyPDF2>=3.0.0",
35
+ "python-docx>=0.8.11",
36
+ "openpyxl>=3.0.0"
37
+ ]
38
+
39
+ [project.optional-dependencies]
40
+ mysql = ["sqlalchemy>=1.4.0", "mysql-connector-python>=8.0.0"]
41
+ postgres = ["sqlalchemy>=1.4.0", "psycopg2-binary>=2.9.0"]
42
+ embeddings = ["sentence-transformers>=2.0.0"]
43
+ all = [
44
+ "sqlalchemy>=1.4.0",
45
+ "mysql-connector-python>=8.0.0",
46
+ "psycopg2-binary>=2.9.0",
47
+ "sentence-transformers>=2.0.0"
48
+ ]
49
+
50
+ [project.urls]
51
+ Homepage = "https://github.com/adityabatta/QuerySUTRA"
52
+ Repository = "https://github.com/adityabatta/QuerySUTRA"
53
+ Issues = "https://github.com/adityabatta/QuerySUTRA/issues"
54
+
55
+ [tool.setuptools.packages.find]
56
+ include = ["sutra*"]
57
+ exclude = ["tests*", "examples*", "data*", "venv*"]
@@ -1,4 +1,4 @@
1
- openai==1.12.0
1
+ openai>=1.0.0
2
2
  pandas==2.2.0
3
3
  numpy==1.24.3
4
4
  matplotlib==3.7.1
@@ -0,0 +1,3 @@
1
+ from setuptools import setup,find_packages
2
+ with open("README.md","r",encoding="utf-8") as f:d=f.read()
3
+ setup(name="QuerySUTRA",version="0.6.1",author="Aditya Batta",description="AI-powered data analysis for structured and unstructured data. Query PDF, Word, CSV, Excel with natural language.",long_description=d,long_description_content_type="text/markdown",packages=find_packages(),python_requires=">=3.8",install_requires=["pandas>=1.3.0","numpy>=1.21.0","openai>=1.0.0","plotly>=5.0.0","matplotlib>=3.3.0","PyPDF2>=3.0.0","python-docx>=0.8.11","openpyxl>=3.0.0"],extras_require={"mysql":["sqlalchemy>=1.4.0","mysql-connector-python>=8.0.0"],"postgres":["sqlalchemy>=1.4.0","psycopg2-binary>=2.9.0"],"embeddings":["sentence-transformers>=2.0.0"],"all":["sqlalchemy>=1.4.0","mysql-connector-python>=8.0.0","psycopg2-binary>=2.9.0","sentence-transformers>=2.0.0"]})
@@ -0,0 +1,6 @@
1
+ """QuerySUTRA v0.6.1 - AI-powered data analysis for structured and unstructured data"""
2
+ __version__ = "0.6.1"
3
+
4
+ from .sutra import SUTRA, QueryResult
5
+
6
+ __all__ = ["SUTRA", "QueryResult", "__version__"]
@@ -1,196 +1,236 @@
1
- """Database management for both SQLite and MySQL"""
2
-
3
- import sqlite3
4
- import pandas as pd
5
- from pathlib import Path
6
- from typing import Optional, Tuple, List
7
- from tabulate import tabulate
8
- import config
9
-
10
- # Add MySQL support
11
- try:
12
- import mysql.connector
13
- MYSQL_AVAILABLE = True
14
- except ImportError:
15
- MYSQL_AVAILABLE = False
16
- print("⚠️ MySQL not installed. Run: pip install mysql-connector-python")
17
-
18
- class DatabaseManager:
19
- """Manage database operations (SQLite or MySQL)"""
20
-
21
- """Database management for both SQLite and MySQL"""
22
-
23
- import sqlite3
24
- import pandas as pd
25
- from pathlib import Path
26
- from typing import Optional, Tuple, List
27
- from tabulate import tabulate
28
- import config
29
-
30
- # Add MySQL support
31
- try:
32
- import mysql.connector
33
- MYSQL_AVAILABLE = True
34
- except ImportError:
35
- MYSQL_AVAILABLE = False
36
- print("⚠️ MySQL not installed. Run: pip install mysql-connector-python")
37
-
38
- class DatabaseManager:
39
- """Manage database operations (SQLite or MySQL)"""
40
-
41
- def __init__(self, db_path: str = ':memory:', db_type: str = 'sqlite'): # FIX: Added indentation
42
- self.db_type = db_type.lower()
43
-
44
- if self.db_type == 'mysql':
45
- if not MYSQL_AVAILABLE:
46
- print("❌ MySQL not available, falling back to SQLite")
47
- self.db_type = 'sqlite'
48
- else:
49
- # First connect without database to create it if needed
50
- try:
51
- conn_temp = mysql.connector.connect(
52
- host=config.MYSQL_HOST,
53
- user=config.MYSQL_USER,
54
- password=config.MYSQL_PASSWORD
55
- )
56
- cursor_temp = conn_temp.cursor()
57
- cursor_temp.execute(f"CREATE DATABASE IF NOT EXISTS {config.MYSQL_DATABASE}")
58
- conn_temp.close()
59
- print(f"✅ Database {config.MYSQL_DATABASE} ready")
60
- except Exception as e:
61
- print(f"❌ Could not create database: {e}")
62
-
63
- # Now connect to the database
64
- self.conn = mysql.connector.connect(
65
- host=config.MYSQL_HOST,
66
- user=config.MYSQL_USER,
67
- password=config.MYSQL_PASSWORD,
68
- database=config.MYSQL_DATABASE
69
- )
70
- self.cursor = self.conn.cursor()
71
- print(f"📂 Connected to MySQL: {config.MYSQL_DATABASE}")
72
-
73
- if self.db_type == 'sqlite': # FIX: Added this block for SQLite
74
- self.conn = sqlite3.connect(db_path)
75
- self.cursor = self.conn.cursor()
76
- print(f"📂 SQLite {'created in memory' if db_path == ':memory:' else f'connected: {db_path}'}")
77
-
78
- # Rest of the methods stay the same...
79
-
80
- def execute_schema(self, schema_sql: str) -> bool:
81
- """Execute SQL schema with MySQL compatibility"""
82
- try:
83
- if self.db_type == 'mysql':
84
- # MySQL adjustments
85
- schema_sql = schema_sql.replace('INTEGER PRIMARY KEY AUTOINCREMENT',
86
- 'INT PRIMARY KEY AUTO_INCREMENT')
87
- schema_sql = schema_sql.replace('TEXT', 'VARCHAR(255)')
88
- schema_sql = schema_sql.replace('REAL', 'DECIMAL(10,2)')
89
-
90
- # Execute statements one by one for MySQL
91
- for statement in schema_sql.split(';'):
92
- if statement.strip():
93
- self.cursor.execute(statement)
94
- self.conn.commit()
95
- else:
96
- # SQLite can handle multiple statements
97
- self.cursor.executescript(schema_sql)
98
- self.conn.commit()
99
-
100
- print("✅ Schema executed successfully!")
101
- return True
102
- except Exception as e:
103
- print(f"❌ Error executing schema: {e}")
104
- return False
105
-
106
- def execute_query(self, query: str) -> Optional[pd.DataFrame]:
107
- """Execute query on either database"""
108
- try:
109
- df = pd.read_sql_query(query, self.conn)
110
- return df
111
- except Exception as e:
112
- print(f"❌ Query error: {e}")
113
- return None
114
-
115
- def get_tables(self):
116
- """Get list of all tables in database"""
117
- if self.db_type == 'mysql':
118
- cursor = self.conn.cursor()
119
- cursor.execute("SHOW TABLES")
120
- tables = [table[0] for table in cursor.fetchall()]
121
- cursor.close()
122
- return tables
123
- else: # sqlite
124
- cursor = self.conn.cursor()
125
- cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
126
- tables = [table[0] for table in cursor.fetchall()]
127
- cursor.close()
128
- return tables
129
-
130
- def get_columns(self, table_name):
131
- """Get list of columns for a specific table"""
132
- if self.db_type == 'mysql':
133
- cursor = self.conn.cursor()
134
- cursor.execute(f"SHOW COLUMNS FROM `{table_name}`")
135
- columns = [col[0] for col in cursor.fetchall()]
136
- cursor.close()
137
- return columns
138
- else: # sqlite
139
- cursor = self.conn.cursor()
140
- cursor.execute(f"PRAGMA table_info({table_name})")
141
- columns = [col[1] for col in cursor.fetchall()]
142
- cursor.close()
143
- return columns
144
-
145
- def get_schema_context(self) -> str:
146
- """Get database schema"""
147
- if self.db_type == 'mysql':
148
- tables = self.get_tables()
149
- schema = []
150
- for table in tables:
151
- self.cursor.execute(f"SHOW CREATE TABLE {table}")
152
- schema.append(self.cursor.fetchone()[1])
153
- return '\n'.join(schema)
154
- else:
155
- self.cursor.execute(
156
- "SELECT sql FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
157
- )
158
- return '\n'.join([row[0] for row in self.cursor.fetchall()])
159
-
160
- def display_tables(self): # FIX: Proper indentation - part of class
161
- """Display all tables with their structure and data"""
162
- tables = self.get_tables()
163
- print(f"\n📋 Created {len(tables)} tables:")
164
-
165
- for table in tables:
166
- print(f"\n Table: {table}")
167
-
168
- # Show columns
169
- columns = self.get_table_info(table)
170
- for col in columns:
171
- print(f" - {col[1]} ({col[2]})")
172
-
173
- # Show row count
174
- count = self.get_row_count(table)
175
- print(f" Records: {count}")
176
-
177
- def get_table_info(self, table_name: str) -> List[Tuple]: # FIX: Proper indentation
178
- """Get column information for a table"""
179
- if self.db_type == 'mysql':
180
- self.cursor.execute(f"DESCRIBE {table_name}")
181
- return [(i, row[0], row[1]) for i, row in enumerate(self.cursor.fetchall())]
182
- else:
183
- self.cursor.execute(f"PRAGMA table_info({table_name})")
184
- return self.cursor.fetchall()
185
-
186
- def get_row_count(self, table_name: str) -> int: # FIX: Proper indentation
187
- """Get number of rows in a table"""
188
- self.cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
189
- return self.cursor.fetchone()[0]
190
-
191
- def close(self):
192
- """Close database connection"""
193
- self.conn.close()
194
- print("📂 Database connection closed")
195
-
1
+ """Database management for both SQLite and MySQL"""
2
+
3
+ import sqlite3
4
+ import pandas as pd
5
+ from pathlib import Path
6
+ from typing import Optional, Tuple, List
7
+ from tabulate import tabulate
8
+ import config
9
+
10
+ # Add MySQL support
11
+ try:
12
+ import mysql.connector
13
+ MYSQL_AVAILABLE = True
14
+ except ImportError:
15
+ MYSQL_AVAILABLE = False
16
+ print("⚠️ MySQL not installed. Run: pip install mysql-connector-python")
17
+
18
+ class DatabaseManager:
19
+ """Manage database operations (SQLite or MySQL)"""
20
+
21
+ """Database management for both SQLite and MySQL"""
22
+
23
+ import sqlite3
24
+ import pandas as pd
25
+ from pathlib import Path
26
+ from typing import Optional, Tuple, List
27
+ from tabulate import tabulate
28
+ import config
29
+
30
+ # Add MySQL support
31
+ try:
32
+ import mysql.connector
33
+ MYSQL_AVAILABLE = True
34
+ except ImportError:
35
+ MYSQL_AVAILABLE = False
36
+ print("⚠️ MySQL not installed. Run: pip install mysql-connector-python")
37
+
38
+ class DatabaseManager:
39
+ """Manage database operations (SQLite or MySQL)"""
40
+
41
+ def __init__(self, db_path: str = ':memory:', db_type: str = 'sqlite'): # FIX: Added indentation
42
+ self.db_type = db_type.lower()
43
+
44
+ if self.db_type == 'mysql':
45
+ if not MYSQL_AVAILABLE:
46
+ print("❌ MySQL not available, falling back to SQLite")
47
+ self.db_type = 'sqlite'
48
+ else:
49
+ # First connect without database to create it if needed
50
+ try:
51
+ conn_temp = mysql.connector.connect(
52
+ host=config.MYSQL_HOST,
53
+ user=config.MYSQL_USER,
54
+ password=config.MYSQL_PASSWORD
55
+ )
56
+ cursor_temp = conn_temp.cursor()
57
+ cursor_temp.execute(f"CREATE DATABASE IF NOT EXISTS {config.MYSQL_DATABASE}")
58
+ conn_temp.close()
59
+ print(f"✅ Database {config.MYSQL_DATABASE} ready")
60
+ except Exception as e:
61
+ print(f"❌ Could not create database: {e}")
62
+
63
+ # Now connect to the database
64
+ self.conn = mysql.connector.connect(
65
+ host=config.MYSQL_HOST,
66
+ user=config.MYSQL_USER,
67
+ password=config.MYSQL_PASSWORD,
68
+ database=config.MYSQL_DATABASE
69
+ )
70
+ self.cursor = self.conn.cursor()
71
+ print(f"📂 Connected to MySQL: {config.MYSQL_DATABASE}")
72
+
73
+ if self.db_type == 'sqlite': # FIX: Added this block for SQLite
74
+ self.conn = sqlite3.connect(db_path)
75
+ self.cursor = self.conn.cursor()
76
+ print(f"📂 SQLite {'created in memory' if db_path == ':memory:' else f'connected: {db_path}'}")
77
+
78
+ # Rest of the methods stay the same...
79
+
80
+ def execute_schema(self, schema_sql: str) -> bool:
81
+ """Execute SQL schema with MySQL compatibility"""
82
+ try:
83
+ if self.db_type == 'mysql':
84
+ # MySQL adjustments
85
+ schema_sql = schema_sql.replace('INTEGER PRIMARY KEY AUTOINCREMENT',
86
+ 'INT PRIMARY KEY AUTO_INCREMENT')
87
+ schema_sql = schema_sql.replace('TEXT', 'VARCHAR(255)')
88
+ schema_sql = schema_sql.replace('REAL', 'DECIMAL(10,2)')
89
+
90
+ # Execute statements one by one for MySQL
91
+ for statement in schema_sql.split(';'):
92
+ if statement.strip():
93
+ self.cursor.execute(statement)
94
+ self.conn.commit()
95
+ else:
96
+ # SQLite can handle multiple statements
97
+ self.cursor.executescript(schema_sql)
98
+ self.conn.commit()
99
+
100
+ print("✅ Schema executed successfully!")
101
+ return True
102
+ except Exception as e:
103
+ print(f"❌ Error executing schema: {e}")
104
+ return False
105
+
106
+ def execute_query(self, query: str) -> Optional[pd.DataFrame]:
107
+ """Execute query on either database"""
108
+ try:
109
+ df = pd.read_sql_query(query, self.conn)
110
+ return df
111
+ except Exception as e:
112
+ print(f"❌ Query error: {e}")
113
+ return None
114
+
115
+ def get_tables(self):
116
+ """Get list of all tables in database"""
117
+ if self.db_type == 'mysql':
118
+ cursor = self.conn.cursor()
119
+ cursor.execute("SHOW TABLES")
120
+ tables = [table[0] for table in cursor.fetchall()]
121
+ cursor.close()
122
+ return tables
123
+ else: # sqlite
124
+ cursor = self.conn.cursor()
125
+ cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
126
+ tables = [table[0] for table in cursor.fetchall()]
127
+ cursor.close()
128
+ return tables
129
+
130
+ def get_columns(self, table_name):
131
+ """Get list of columns for a specific table"""
132
+ if self.db_type == 'mysql':
133
+ cursor = self.conn.cursor()
134
+ cursor.execute(f"SHOW COLUMNS FROM `{table_name}`")
135
+ columns = [col[0] for col in cursor.fetchall()]
136
+ cursor.close()
137
+ return columns
138
+ else: # sqlite
139
+ cursor = self.conn.cursor()
140
+ cursor.execute(f"PRAGMA table_info({table_name})")
141
+ columns = [col[1] for col in cursor.fetchall()]
142
+ cursor.close()
143
+ return columns
144
+
145
+ def get_schema_context(self) -> str:
146
+ """Get database schema with relationship information"""
147
+ if self.db_type == 'mysql':
148
+ tables = self.get_tables()
149
+ schema = []
150
+ for table in tables:
151
+ self.cursor.execute(f"SHOW CREATE TABLE {table}")
152
+ schema.append(self.cursor.fetchone()[1])
153
+ schema_text = '\n'.join(schema)
154
+ else:
155
+ self.cursor.execute(
156
+ "SELECT sql FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
157
+ )
158
+ schema_text = '\n'.join([row[0] for row in self.cursor.fetchall()])
159
+
160
+ # Add relationship summary at the beginning
161
+ relationships = self._extract_relationships(schema_text)
162
+ if relationships:
163
+ relationship_summary = "\n=== TABLE RELATIONSHIPS ===\n" + "\n".join(relationships) + "\n\n=== FULL SCHEMA ===\n"
164
+ return relationship_summary + schema_text
165
+ return schema_text
166
+
167
+ def _extract_relationships(self, schema_text: str) -> List[str]:
168
+ """Extract and format foreign key relationships from schema"""
169
+ import re
170
+ relationships = []
171
+
172
+ # Pattern to match FOREIGN KEY statements
173
+ fk_pattern = r'FOREIGN KEY\s*\(([^)]+)\)\s*REFERENCES\s+(\w+)\s*\(([^)]+)\)'
174
+
175
+ # Split schema into individual table definitions
176
+ tables = schema_text.split('CREATE TABLE')
177
+
178
+ for table_def in tables:
179
+ if not table_def.strip():
180
+ continue
181
+
182
+ # Extract table name
183
+ table_match = re.search(r'[`"]?(\w+)[`"]?', table_def)
184
+ if not table_match:
185
+ continue
186
+ table_name = table_match.group(1)
187
+
188
+ # Find all foreign keys in this table
189
+ for match in re.finditer(fk_pattern, table_def, re.IGNORECASE):
190
+ fk_column = match.group(1).strip('`" ')
191
+ ref_table = match.group(2).strip('`" ')
192
+ ref_column = match.group(3).strip('`" ')
193
+
194
+ relationships.append(
195
+ f" {table_name}.{fk_column} → {ref_table}.{ref_column}"
196
+ )
197
+
198
+ return relationships
199
+
200
+ def display_tables(self): # FIX: Proper indentation - part of class
201
+ """Display all tables with their structure and data"""
202
+ tables = self.get_tables()
203
+ print(f"\n📋 Created {len(tables)} tables:")
204
+
205
+ for table in tables:
206
+ print(f"\n Table: {table}")
207
+
208
+ # Show columns
209
+ columns = self.get_table_info(table)
210
+ for col in columns:
211
+ print(f" - {col[1]} ({col[2]})")
212
+
213
+ # Show row count
214
+ count = self.get_row_count(table)
215
+ print(f" Records: {count}")
216
+
217
+ def get_table_info(self, table_name: str) -> List[Tuple]: # FIX: Proper indentation
218
+ """Get column information for a table"""
219
+ if self.db_type == 'mysql':
220
+ self.cursor.execute(f"DESCRIBE {table_name}")
221
+ return [(i, row[0], row[1]) for i, row in enumerate(self.cursor.fetchall())]
222
+ else:
223
+ self.cursor.execute(f"PRAGMA table_info({table_name})")
224
+ return self.cursor.fetchall()
225
+
226
+ def get_row_count(self, table_name: str) -> int: # FIX: Proper indentation
227
+ """Get number of rows in a table"""
228
+ self.cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
229
+ return self.cursor.fetchone()[0]
230
+
231
+ def close(self):
232
+ """Close database connection"""
233
+ self.conn.close()
234
+ print("📂 Database connection closed")
235
+
196
236