richardr1126 commited on
Commit
485ca65
1 Parent(s): 2e9310d
Files changed (1) hide show
  1. app.py +12 -1
app.py CHANGED
@@ -132,9 +132,20 @@ def extract_db_code(text):
132
  return [match.strip() for match in matches]
133
 
134
  def extract_from_code_block(text):
 
135
  pattern = r'```(?:\w+)?\s?(.*?)```'
136
  match = re.search(pattern, text, re.DOTALL)
137
- return match.group(1).strip() if match else ''
 
 
 
 
 
 
 
 
 
 
138
 
139
  def generate_dummy_db(db_info, question):
140
  pre_prompt = """
 
132
  return [match.strip() for match in matches]
133
 
134
  def extract_from_code_block(text):
135
+ # First try to match text inside ```
136
  pattern = r'```(?:\w+)?\s?(.*?)```'
137
  match = re.search(pattern, text, re.DOTALL)
138
+ if match:
139
+ return match.group(1).strip()
140
+
141
+ # If no match in ```, try to match SQL-like text
142
+ pattern = r'\((SELECT .*?)\)'
143
+ match = re.search(pattern, text, re.DOTALL)
144
+ if match:
145
+ return match.group(1).strip()
146
+
147
+ # Return None if no match found
148
+ return ""
149
 
150
  def generate_dummy_db(db_info, question):
151
  pre_prompt = """