resolve dependency

Author: Chenglong-MS

py-src/data_formulator/agents/agent_query_completion.py

@@ -54,6 +54,10 @@ def __init__(self, client):
 
     def run(self, data_source_metadata, query):
 
+        # For MongoDB, treat it as a SQL-like data source for query generation
+        if data_source_metadata['data_loader_type'] == "mongodb":
+            data_source_metadata['data_loader_type'] = "SQL"
+
         user_query = f"[DATA SOURCE]\n\n{json.dumps(data_source_metadata, indent=2)}\n\n[USER INPUTS]\n\n{query}\n\n"
 
         logger.info(user_query)

py-src/data_formulator/data_loader/__init__.py

@@ -5,6 +5,7 @@
 from data_formulator.data_loader.s3_data_loader import S3DataLoader
 from data_formulator.data_loader.azure_blob_data_loader import AzureBlobDataLoader
 from data_formulator.data_loader.postgresql_data_loader import PostgreSQLDataLoader
+from data_formulator.data_loader.mongodb_data_loader import MongoDBDataLoader
 from data_formulator.data_loader.bigquery_data_loader import BigQueryDataLoader
 
 DATA_LOADERS = {
@@ -14,8 +15,8 @@
     "s3": S3DataLoader,
     "azure_blob": AzureBlobDataLoader,
     "postgresql": PostgreSQLDataLoader,
+    "mongodb": MongoDBDataLoader,
     "bigquery": BigQueryDataLoader
 }
 
-__all__ = ["ExternalDataLoader", "MySQLDataLoader", "MSSQLDataLoader", "KustoDataLoader", "S3DataLoader", "AzureBlobDataLoader","PostgreSQLDataLoader", "BigQueryDataLoader", "DATA_LOADERS"]
-
+__all__ = ["ExternalDataLoader", "MySQLDataLoader", "MSSQLDataLoader", "KustoDataLoader", "S3DataLoader", "AzureBlobDataLoader","PostgreSQLDataLoader", "MongoDBDataLoader", "BigQueryDataLoader", "DATA_LOADERS"]