Added Models Support

app.py

@@ -15,7 +15,15 @@ print('Connecting to DB...')
 # Connect to DB
 conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)
 
-llm_engine = HfEngine(model="meta-llama/Meta-Llama-3-70B-Instruct")
+models = ["meta-llama/Meta-Llama-3-70B-Instruct",
+          "meta-llama/Llama-3.1-70B-Instruct"]
+
+for model in models:
+  llm_engine = HfEngine(model=model)
+  model_status = llm_engine.client.get_model_status()
+
+  if model_status.loaded:
+    break
 
 def get_schemas():
     schemas = conn.execute("""
@@ -90,8 +98,7 @@ class SQLExecutorTool(Tool):
     output_type = "pandas.core.frame.DataFrame"
 
     def forward(self, query: str) -> str:
-        with duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True) as con:
-            output_df = conn.sql(query).df()
+        output_df = conn.sql(query).df()
         return output_df
     
 tool = SQLExecutorTool()