feat(rag): Support RAG SDK (#1322)

dbgpt/datasource/__init__.py

@@ -1,6 +1,17 @@
 """Module to define the data source connectors."""
 
+from typing import Any
+
 from .base import BaseConnector  # noqa: F401
-from .rdbms.base import RDBMSConnector  # noqa: F401
+
+
+def __getattr__(name: str) -> Any:
+    if name == "RDBMSConnector":
+        from .rdbms.base import RDBMSConnector  # noqa: F401
+
+        return RDBMSConnector
+    else:
+        raise AttributeError(f"Could not find: {name} in datasource")
+
 
 __ALL__ = ["BaseConnector", "RDBMSConnector"]

dbgpt/datasource/operators/datasource_operator.py

@@ -12,10 +12,10 @@ from ..base import BaseConnector
 class DatasourceOperator(MapOperator[str, Any]):
     """The Datasource Operator."""
 
-    def __init__(self, connection: BaseConnector, **kwargs):
+    def __init__(self, connector: BaseConnector, **kwargs):
         """Create the datasource operator."""
         super().__init__(**kwargs)
-        self._connection = connection
+        self._connector = connector
 
     async def map(self, input_value: str) -> Any:
         """Execute the query."""
@@ -23,4 +23,4 @@ class DatasourceOperator(MapOperator[str, Any]):
 
     def query(self, input_value: str) -> Any:
         """Execute the query."""
-        return self._connection.run_to_df(input_value)
+        return self._connector.run_to_df(input_value)

dbgpt/datasource/rdbms/base.py

@@ -3,11 +3,11 @@
 from __future__ import annotations
 
 import logging
+import re
 from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, cast
 from urllib.parse import quote
 from urllib.parse import quote_plus as urlquote
 
-import regex as re
 import sqlalchemy
 import sqlparse
 from sqlalchemy import MetaData, Table, create_engine, inspect, select, text