Qury checking

app.py

@@ -5,11 +5,32 @@ from langchain.chat_models import ChatOpenAI
 from langchain.prompts.prompt import PromptTemplate
 
 
+def check_query(query):
+    if query.startswith("### Query"):
+        split = query.split('\n\n')
+        q_text = split[0]
+        t_text = split[1]
+
+        if t_text.startswith("### Tables"):
+            tables = t_text.split('\n')
+            print(tables[1:])
+            return('big')
+        else:
+            return('error')
+    return('small')
+
 def answer_question(query):
-  if query.startswith("### Query"):
-     split = query.split('\n\n')
-     print(split)
-  return("DONE: " + query)
+    query_check = check_query(query)
+    if query_check == 'big':
+        return("BIG TABLE")
+    if query_check == 'small':
+        return('SMALL TABLE')
+    if query_check == 'error':
+        return('ERROR: Wrong format for getting the big db schema')
+  
+
+
+    return("DONE: " + query)
 
 if __name__ == "__main__":
     import gradio as gr