github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-16 06:30:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: Refactor datasource module (#1309)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-03-18 18:06:40 +08:00
committed by GitHub
parent 84bedee306
commit 4970c9f813
108 changed files with 1194 additions and 1066 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/awel/simple_nl_schema_sql_chart_example.py
View File

@@ -7,8 +7,8 @@ from pydantic import BaseModel, Field
from dbgpt.configs.model_config import MODEL_PATH, PILOT_PATH
from dbgpt.core import LLMClient, ModelMessage, ModelMessageRoleType, ModelRequest
from dbgpt.core.awel import DAG, HttpTrigger, JoinOperator, MapOperator
from dbgpt.datasource.rdbms.base import RDBMSDatabase
from dbgpt.datasource.rdbms.conn_sqlite import SQLiteTempConnect
from dbgpt.datasource.rdbms.base import RDBMSConnector
from dbgpt.datasource.rdbms.conn_sqlite import SQLiteTempConnector
from dbgpt.model.proxy import OpenAILLMClient
from dbgpt.rag.embedding.embedding_factory import DefaultEmbeddingFactory
from dbgpt.rag.operators.schema_linking import SchemaLinkingOperator
@@ -66,7 +66,7 @@ def _create_vector_connector():
def _create_temporary_connection():
"""Create a temporary database connection for testing."""
connect = SQLiteTempConnect.create_temporary_db()
connect = SQLiteTempConnector.create_temporary_db()
connect.create_temp_tables(
{
"user": {
@@ -181,10 +181,10 @@ class SqlGenOperator(MapOperator[Any, Any]):
class SqlExecOperator(MapOperator[Any, Any]):
"""The Sql Execution Operator."""
def __init__(self, connection: Optional[RDBMSDatabase] = None, **kwargs):
def __init__(self, connection: Optional[RDBMSConnector] = None, **kwargs):
"""
Args:
connection (Optional[RDBMSDatabase]): RDBMSDatabase connection
connection (Optional[RDBMSConnector]): RDBMSConnector connection
"""
super().__init__(**kwargs)
self._connection = connection
@@ -207,7 +207,7 @@ class ChartDrawOperator(MapOperator[Any, Any]):
def __init__(self, **kwargs):
"""
Args:
connection (RDBMSDatabase): The connection.
connection (RDBMSConnector): The connection.
"""
super().__init__(**kwargs)