CD17 commited on
Commit
906657d
·
verified ·
1 Parent(s): fb6e1bb

Upload functions_llama.py

Browse files
Files changed (1) hide show
  1. functions_llama.py +310 -0
functions_llama.py ADDED
@@ -0,0 +1,310 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain_community.utilities import SQLDatabase
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+ from langchain_core.prompts import ChatPromptTemplate
5
+ from langchain_core.output_parsers import StrOutputParser
6
+ from langchain_core.runnables import RunnablePassthrough
7
+ import re
8
+ from typing import Dict, List
9
+
10
+ # Load model and tokenizer
11
+ model_name = "meta-llama/Llama-3.3-70B-Instruct"
12
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
13
+ model = AutoModelForCausalLM.from_pretrained(model_name)
14
+
15
+ class HuggingFaceLLM:
16
+ def __init__(self, model, tokenizer):
17
+ self.model = model
18
+ self.tokenizer = tokenizer
19
+
20
+ def __call__(self, prompt):
21
+ inputs = self.tokenizer(prompt, return_tensors="pt")
22
+ with torch.no_grad():
23
+ outputs = self.model.generate(**inputs, max_length=512)
24
+ return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
25
+
26
+ # Initialize LLM
27
+ llm = HuggingFaceLLM(model, tokenizer)
28
+
29
+ # Initialize database connection
30
+ db = SQLDatabase.from_uri("sqlite:///Spring_2025_courses.db", sample_rows_in_table_info=0)
31
+
32
+ # Helper functions
33
+ def get_schema(_):
34
+ """Retrieve database schema"""
35
+ return db.get_table_info()
36
+
37
+ def run_query(query):
38
+ """Execute SQL query"""
39
+ return db.run(query)
40
+
41
+ # SQL generation prompt
42
+ system_prompt = """
43
+ You are a SQLite expert. Given an input question, create one syntactically correct SQLite query to run. Generate only one query. No pre-amble.
44
+
45
+ Here is the relevant table information:
46
+ schema: {schema}
47
+ Departments: 'Applied Math & Statistics', 'ART', 'Biomedical Science (BMS)', 'Computer Science (CS), Computer Networks and Cybersecurity (CNCS)', 'Data Science (DS)', 'Dance', 'General Education (GE)'
48
+ New/Repeat: Have no values (Don't use New/Repeat in queries)
49
+
50
+ Tips:
51
+ - Use LIKE instead of = in the queries
52
+
53
+ Below are a number of examples of questions and their corresponding SQL queries.
54
+
55
+ example 1:
56
+ Question:List all departments.
57
+ SQL query:SELECT Department FROM Spring_2025_courses;
58
+
59
+ example 2:
60
+ Question:Find all classes that are taught by Dr. Qu.
61
+ SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Instructor LIKE '%Qu%';
62
+
63
+ example 3:
64
+ Question:List all courses in the Data Science department.
65
+ SQL query:SELECT CourseCode, CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Data Science%';
66
+
67
+ example 4:
68
+ Question:Find the total number of courses offered for Spring 2025.
69
+ SQL query:SELECT COUNT(CourseCode) FROM Spring_2025_courses;
70
+
71
+ example 5:
72
+ Question:List all professors from Biomed Department.
73
+ SQL query:SELECT Instructor FROM Spring_2025_courses WHERE Department LIKE '%Biomed%';
74
+
75
+ example 6:
76
+ Question:Which courses are available in the Biomed department for Spring 2025?
77
+ SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Biomed%';
78
+
79
+ example 7:
80
+ Question:How many courses are offered under the Arts Management (BFA) program?
81
+ SQL query:SELECT COUNT(CourseCode) FROM Spring_2025_courses WHERE Department LIKE Program LIKE '%BFA%'
82
+
83
+ example 8:
84
+ Question:Which courses are available for the Master of Science in Applied Math?
85
+ SQL query:SELECT CourseTitle FROM Spring_2025_courses WHERE Program LIKE '%Master of Science%' AND Department LIKE '%Applied Math%';
86
+
87
+ Write only one SQLite query that would answer the user's question. No pre-amble.
88
+ """
89
+
90
+ human_prompt = """Based on the table schema below, write a SQL query that would answer the user's question:
91
+ {schema}
92
+
93
+ Question: {question}
94
+ SQL Query:"""
95
+
96
+ prompt = ChatPromptTemplate.from_messages([
97
+ ("system", system_prompt),
98
+ ("human", human_prompt),
99
+ ])
100
+
101
+ # Build query generation chain
102
+ sql_generator = (
103
+ RunnablePassthrough.assign(schema=get_schema)
104
+ | prompt
105
+ | llm.bind(stop=["\nSQLResult:"])
106
+ | StrOutputParser()
107
+ )
108
+
109
+ # Natural language response generation
110
+ nl_system_prompt = """Given an input question and SQL response, convert it to a natural language answer.
111
+
112
+ Below are a number of examples of questions and their corresponding SQL queries.
113
+
114
+ example 1:
115
+ Question: Which courses are available in the Data Science department?
116
+ SQL Query: SELECT CourseTitle FROM Spring_2025_courses WHERE Department LIKE '%Data Science%';
117
+ SQL Response: [('Introduction to Data Science',), ('Career Development in Data Science',), ('Data Structures',), ('Data Visualization',), ('Data Inference',), ('Data Mining',), ('Big Data Engineering',), ('Independent Study for Data Science',), ('Senior Project',), ('Mathematical Foundation for Data Science (Online)',), ('Computational Foundation for Data Science (Online)',), ('Probability for Data Science (Online)',), ('Exploratory Data Analysis and Visualization',), ('Modern Applied Statistical Learning',), ('Data Mining for Business',), ('Big Data and Data Engineering',), ('Design and Analysis of Experiments',), ('Design and Analysis of Experiments',), ('Computer Vision and Natural Language Processing (Independent Study)',), ('Capstone Project',)]
118
+ Response: The available courses in the Data Science department are Introduction to Data Science, Career Development in Data Science, Data Structures, Data Visualization, Data Inference, Data Mining, Big Data Engineering, Independent Study for Data Science, Senior Project, Mathematical Foundation for Data Science, Computational Foundation for Data Science, Probability for Data Science, Exploratory Data Analysis and Visualization, Modern Applied Statistical Learning, Data Mining for Business, Big Data and Data Engineering, Design and Analysis of Experiments, and Capstone Project.
119
+
120
+ example 2:
121
+ Question: Find all classes that are taught by Dr. Qu.
122
+ SQL Query: SELECT CourseTitle FROM Spring_2025_courses WHERE Instructor LIKE '%Qu%';
123
+ SQL Response: [('Introduction to Data Science',), ('Data Structures',), ('Data Mining',), ('Computational Foundation for Data Science (Online)',), ('Data Mining for Business',), ('Computer Vision and Natural Language Processing (Independent Study)',)]
124
+ Response: Dr. Qu is teaching Introduction to Data Science, Data Structures, Data Mining, Computational Foundation for Data Science, Data Mining for Business, and Computer Vision and Natural Language Processing.
125
+
126
+ Generate only ONE answer based on the SQL response. No pre-amble.
127
+ """
128
+
129
+ response_prompt = """Based on the table schema below, question, sql query, and sql response, write a natural language response. No pre-amble:
130
+ {schema}
131
+
132
+ Question:{question}
133
+ SQL Query:{query}
134
+ SQL Response:{response}
135
+ Response:"""
136
+
137
+ nl_prompt = ChatPromptTemplate.from_messages([
138
+ ("system", nl_system_prompt),
139
+ ("human", response_prompt),
140
+ ])
141
+
142
+ # Complete chain with natural language output
143
+ complete_chain = (
144
+ RunnablePassthrough.assign(query=sql_generator)
145
+ | RunnablePassthrough.assign(
146
+ schema=get_schema,
147
+ response=lambda x: db.run(x['query']),
148
+ )
149
+ | nl_prompt
150
+ | llm
151
+ )
152
+
153
+ class InputValidator:
154
+ """Input validation and sanitization for database queries"""
155
+
156
+ RESTRICTED_WORDS = {
157
+ 'delete', 'drop', 'truncate', 'update', 'insert',
158
+ 'alter', 'create', 'replace', 'modify', 'grant'
159
+ }
160
+
161
+ ALLOWED_TABLES = ['spring_2025_courses']
162
+
163
+ def __init__(self):
164
+ self.error_messages = []
165
+
166
+ def validate_question(self, question: str) -> bool:
167
+ """
168
+ Validate a natural language question
169
+
170
+ Args:
171
+ question (str): User's input question
172
+
173
+ Returns:
174
+ bool: True if valid, False otherwise
175
+ """
176
+ self.error_messages = []
177
+
178
+ # Check if question is empty or too long
179
+ if not question or not question.strip():
180
+ self.error_messages.append("Question cannot be empty")
181
+ return False
182
+
183
+ if len(question) > 500:
184
+ self.error_messages.append("Question is too long (max 500 characters)")
185
+ return False
186
+
187
+ # Check for basic SQL injection attempts
188
+ question_lower = question.lower()
189
+ if any(word in question_lower for word in self.RESTRICTED_WORDS):
190
+ self.error_messages.append("Question contains restricted keywords")
191
+ return False
192
+
193
+ # Check for excessive special characters
194
+ if re.search(r'[;{}\\]', question):
195
+ self.error_messages.append("Question contains invalid characters")
196
+ return False
197
+
198
+ return True
199
+
200
+ def get_error_messages(self) -> List[str]:
201
+ """Get all error messages from validation"""
202
+ return self.error_messages
203
+
204
+ class QueryValidator:
205
+ """Validate generated SQL queries before execution"""
206
+
207
+ def __init__(self, db_connection):
208
+ self.db = db_connection
209
+ self.error_messages = []
210
+
211
+ def get_error_messages(self) -> List[str]:
212
+ """Get all error messages from validation"""
213
+ return self.error_messages
214
+
215
+ def validate_sql(self, sql: str) -> bool:
216
+ """Validate generated SQL query"""
217
+ sql_lower = sql.lower()
218
+
219
+ # Check if query is read-only (SELECT only)
220
+ if not sql_lower.strip().startswith('select'):
221
+ self.error_messages.append("Only SELECT queries are allowed")
222
+ return False
223
+
224
+ # Check for multiple statements
225
+ if ';' in sql[:-1]: # Allow semicolon at the end
226
+ self.error_messages.append("Multiple SQL statements are not allowed")
227
+ return False
228
+
229
+ # Validate table names
230
+ table = self._extract_table_names(sql_lower)
231
+ if table not in InputValidator.ALLOWED_TABLES:
232
+ self.error_messages.append("Query contains invalid table names")
233
+ return False
234
+
235
+ return True
236
+
237
+ def _extract_table_names(self, sql: str) -> str:
238
+ """Extract table names from SQL query"""
239
+ from_matches = re.findall(r'from\s+([a-zA-Z_][a-zA-Z0-9_]*)', sql.lower())
240
+ return from_matches[0]
241
+
242
+ class SafeQueryExecutor:
243
+ """Safe execution of natural language queries"""
244
+
245
+ def __init__(self, db_connection, llm_chain):
246
+ self.input_validator = InputValidator()
247
+ self.query_validator = QueryValidator(db_connection)
248
+ self.db = db_connection
249
+ self.chain = llm_chain
250
+
251
+ def execute_safe_query(self, question: str) -> Dict:
252
+ """
253
+ Safely execute a natural language query
254
+
255
+ Args:
256
+ question (str): User's natural language question
257
+
258
+ Returns:
259
+ dict: Query result and status
260
+ """
261
+ # Validate input
262
+ if not self.input_validator.validate_question(question):
263
+ return {
264
+ 'success': False,
265
+ 'error': self.input_validator.get_error_messages(),
266
+ 'query': None,
267
+ 'result': None
268
+ }
269
+
270
+ try:
271
+ # Generate SQL query
272
+ sql_query = sql_generator.invoke({"question": question})
273
+
274
+ # Validate generated SQL
275
+ if not self.query_validator.validate_sql(sql_query):
276
+ return {
277
+ 'success': False,
278
+ 'error': self.query_validator.get_error_messages(),
279
+ 'query': sql_query,
280
+ 'result': None
281
+ }
282
+
283
+ # Execute query
284
+ result = complete_chain.invoke({"question": question})
285
+
286
+ return {
287
+ 'success': True,
288
+ 'error': None,
289
+ 'query': sql_query,
290
+ 'result': result
291
+ }
292
+
293
+ except Exception as e:
294
+ return {
295
+ 'success': False,
296
+ 'error': [str(e)],
297
+ 'query': None,
298
+ 'result': None
299
+ }
300
+
301
+ # Initialize the safe executor
302
+ safe_executor = SafeQueryExecutor(db, complete_chain)
303
+
304
+ # Example queries
305
+ example_queries = [
306
+ "Which courses are available in the Data Science department for Spring 2025?",
307
+ "What are the details of the course 'Introduction to Python'?",
308
+ "How many courses are offered under the Computer Science (Bachelor of Science) program?",
309
+ "Show me the courses taught by Dr. Qu in Spring 2025."
310
+ ]