feat(datasource):add oceanbase support (#1622)

Co-authored-by: csunny <cfqsunny@163.com>
Co-authored-by: aries_ckt <916701291@qq.com>

dbgpt/app/base.py

@@ -8,6 +8,7 @@ from typing import Optional
 
 from dbgpt._private.config import Config
 from dbgpt.component import SystemApp
+from dbgpt.storage import DBType
 from dbgpt.util.parameter_utils import BaseParameters
 
 ROOT_PATH = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -31,7 +32,6 @@ def async_db_summary(system_app: SystemApp):
 
 
 def server_init(param: "WebServerParameters", system_app: SystemApp):
-
     # logger.info(f"args: {args}")
     # init config
     cfg = Config()
@@ -105,13 +105,16 @@ def _initialize_db(
     from urllib.parse import quote_plus as urlquote
 
     from dbgpt.configs.model_config import PILOT_PATH
+    from dbgpt.datasource.rdbms.dialect.oceanbase.ob_dialect import (  # noqa: F401
+        OBDialect,
+    )
     from dbgpt.storage.metadata.db_manager import initialize_db
 
     CFG = Config()
     db_name = CFG.LOCAL_DB_NAME
     default_meta_data_path = os.path.join(PILOT_PATH, "meta_data")
     os.makedirs(default_meta_data_path, exist_ok=True)
-    if CFG.LOCAL_DB_TYPE == "mysql":
+    if CFG.LOCAL_DB_TYPE == DBType.MySQL.value():
         db_url = (
             f"mysql+pymysql://{quote(CFG.LOCAL_DB_USER)}:"
             f"{urlquote(CFG.LOCAL_DB_PASSWORD)}@"
@@ -121,6 +124,15 @@ def _initialize_db(
         )
         # Try to create database, if failed, will raise exception
         _create_mysql_database(db_name, db_url, try_to_create_db)
+    elif CFG.LOCAL_DB_TYPE == DBType.OceanBase.value():
+        db_url = (
+            f"mysql+ob://{quote(CFG.LOCAL_DB_USER)}:"
+            f"{urlquote(CFG.LOCAL_DB_PASSWORD)}@"
+            f"{CFG.LOCAL_DB_HOST}:"
+            f"{str(CFG.LOCAL_DB_PORT)}/"
+            f"{db_name}?charset=utf8mb4"
+        )
+        _create_mysql_database(db_name, db_url, try_to_create_db)
     else:
         sqlite_db_path = os.path.join(default_meta_data_path, f"{db_name}.db")
         db_url = f"sqlite:///{sqlite_db_path}"

dbgpt/app/static/_next/static/LQoR86Gp1MW_leEBmkfhZ/_buildManifest.js

@@ -1 +1 @@
-self.__BUILD_MANIFEST=function(s,c,a,e,t,n,f,d,k,h,i,u,b,j,p,g,o,l,r,_){return{__rewrites:{beforeFiles:[],afterFiles:[],fallback:[]},"/":[p,s,c,e,a,f,h,d,g,"static/chunks/9305-f44429d5185a9fc7.js","static/chunks/1353-705aa47cc2b94999.js","static/chunks/pages/index-0b2d61c1c6358f20.js"],"/_error":["static/chunks/pages/_error-8095ba9e1bf12f30.js"],"/agent":[s,c,a,t,h,n,"static/chunks/pages/agent-2be7990da37f5165.js"],"/app":[i,s,c,e,a,t,n,u,b,o,j,"static/chunks/pages/app-8154f6fcced2f743.js"],"/chat":["static/chunks/pages/chat-5794854c0948b84c.js"],"/database":[s,c,e,a,t,n,d,k,"static/chunks/3718-e111d727d432bdd2.js","static/chunks/pages/database-7384ab94b08f23ff.js"],"/flow":[i,s,c,a,u,b,o,j,"static/chunks/pages/flow-33fe9f396642fb4c.js"],"/flow/canvas":[p,i,s,c,e,a,f,d,u,k,b,l,g,"static/chunks/1425-6e94ae18b1ac5a70.js",j,"static/chunks/pages/flow/canvas-644b6ee718585173.js"],"/knowledge":[r,s,c,e,a,t,h,n,d,k,_,l,"static/chunks/5237-1d36a3742424b75e.js","static/chunks/pages/knowledge-223d50e9531bd961.js"],"/knowledge/chunk":[s,e,t,f,n,"static/chunks/pages/knowledge/chunk-625a32aed5f380e2.js"],"/knowledge/graph":["static/chunks/90912e1b-ed32608ee46ab40f.js","static/chunks/193-5e83ce3fd4f165ef.js","static/chunks/pages/knowledge/graph-9fb1ec6bf06d5108.js"],"/models":[r,s,c,e,a,k,"static/chunks/3444-30181eacc7980e66.js","static/chunks/pages/models-446238c56e41aa1b.js"],"/prompt":[s,c,e,a,f,_,"static/chunks/7184-3ca3f58327a6986a.js","static/chunks/7869-1a99e25b182b3eaa.js","static/chunks/pages/prompt-c44ac718b4d637c9.js"],sortedPages:["/","/_app","/_error","/agent","/app","/chat","/database","/flow","/flow/canvas","/knowledge","/knowledge/chunk","/knowledge/graph","/models","/prompt"]}}("static/chunks/2185-30f9d0578fa0d631.js","static/chunks/5503-c65f6d730754acc7.js","static/chunks/9479-21f588e1fd4e6b6d.js","static/chunks/1009-4b2af86bde623424.js","static/chunks/785-c3544abc036fc97d.js","static/chunks/5813-c6244a8eba7ef4ae.js","static/chunks/1647-8683da4db89d68c1.js","static/chunks/411-b5d3e7f64bee2335.js","static/chunks/8928-0e78def492052d13.js","static/chunks/4553-61740188e6a650a8.js","static/chunks/971df74e-7436ff4085ebb785.js","static/chunks/7434-29506257e67e8077.js","static/chunks/9924-5bce555f07385e1f.js","static/css/b4846eed11c4725f.css","static/chunks/29107295-75edf0bf34e24b1e.js","static/chunks/2487-cda9d2a2fd712a15.js","static/chunks/6165-93d23bc520382b2c.js","static/chunks/2282-96412afca1591c9a.js","static/chunks/75fc9c18-1d6133135d3d283c.js","static/chunks/5733-7ef320ab0f876a5e.js"),self.__BUILD_MANIFEST_CB&&self.__BUILD_MANIFEST_CB();
+self.__BUILD_MANIFEST=function(s,c,a,e,t,n,f,d,k,h,i,u,b,j,p,g,o,l,r,_){return{__rewrites:{beforeFiles:[],afterFiles:[],fallback:[]},"/":[p,s,c,e,a,f,h,d,g,"static/chunks/9305-f44429d5185a9fc7.js","static/chunks/1353-705aa47cc2b94999.js","static/chunks/pages/index-11135322ea50180f.js"],"/_error":["static/chunks/pages/_error-8095ba9e1bf12f30.js"],"/agent":[s,c,a,t,h,n,"static/chunks/pages/agent-2be7990da37f5165.js"],"/app":[i,s,c,e,a,t,n,u,b,o,j,"static/chunks/pages/app-8154f6fcced2f743.js"],"/chat":["static/chunks/pages/chat-5794854c0948b84c.js"],"/database":[s,c,e,a,t,n,d,k,"static/chunks/3718-e111d727d432bdd2.js","static/chunks/pages/database-7384ab94b08f23ff.js"],"/flow":[i,s,c,a,u,b,o,j,"static/chunks/pages/flow-33fe9f396642fb4c.js"],"/flow/canvas":[p,i,s,c,e,a,f,d,u,k,b,l,g,"static/chunks/1425-6e94ae18b1ac5a70.js",j,"static/chunks/pages/flow/canvas-644b6ee718585173.js"],"/knowledge":[r,s,c,e,a,t,h,n,d,k,_,l,"static/chunks/5237-1d36a3742424b75e.js","static/chunks/pages/knowledge-223d50e9531bd961.js"],"/knowledge/chunk":[s,e,t,f,n,"static/chunks/pages/knowledge/chunk-625a32aed5f380e2.js"],"/knowledge/graph":["static/chunks/90912e1b-ed32608ee46ab40f.js","static/chunks/193-5e83ce3fd4f165ef.js","static/chunks/pages/knowledge/graph-9fb1ec6bf06d5108.js"],"/models":[r,s,c,e,a,k,"static/chunks/3444-30181eacc7980e66.js","static/chunks/pages/models-446238c56e41aa1b.js"],"/prompt":[s,c,e,a,f,_,"static/chunks/7184-3ca3f58327a6986a.js","static/chunks/7869-1a99e25b182b3eaa.js","static/chunks/pages/prompt-c44ac718b4d637c9.js"],sortedPages:["/","/_app","/_error","/agent","/app","/chat","/database","/flow","/flow/canvas","/knowledge","/knowledge/chunk","/knowledge/graph","/models","/prompt"]}}("static/chunks/2185-30f9d0578fa0d631.js","static/chunks/5503-c65f6d730754acc7.js","static/chunks/9479-21f588e1fd4e6b6d.js","static/chunks/1009-4b2af86bde623424.js","static/chunks/785-c3544abc036fc97d.js","static/chunks/5813-c6244a8eba7ef4ae.js","static/chunks/1647-8683da4db89d68c1.js","static/chunks/411-b5d3e7f64bee2335.js","static/chunks/8928-0e78def492052d13.js","static/chunks/4553-61740188e6a650a8.js","static/chunks/971df74e-7436ff4085ebb785.js","static/chunks/7434-29506257e67e8077.js","static/chunks/9924-5bce555f07385e1f.js","static/css/b4846eed11c4725f.css","static/chunks/29107295-75edf0bf34e24b1e.js","static/chunks/2487-cda9d2a2fd712a15.js","static/chunks/6165-93d23bc520382b2c.js","static/chunks/2282-96412afca1591c9a.js","static/chunks/75fc9c18-1d6133135d3d283c.js","static/chunks/5733-7ef320ab0f876a5e.js"),self.__BUILD_MANIFEST_CB&&self.__BUILD_MANIFEST_CB();

dbgpt/datasource/manages/connector_manager.py

@@ -49,6 +49,7 @@ class ConnectorManager(BaseComponent):
         from dbgpt.datasource.rdbms.conn_hive import HiveConnector  # noqa: F401
         from dbgpt.datasource.rdbms.conn_mssql import MSSQLConnector  # noqa: F401
         from dbgpt.datasource.rdbms.conn_mysql import MySQLConnector  # noqa: F401
+        from dbgpt.datasource.rdbms.conn_oceanbase import OceanBaseConnect  # noqa: F401
         from dbgpt.datasource.rdbms.conn_postgresql import (  # noqa: F401
             PostgreSQLConnector,
         )
@@ -57,6 +58,9 @@ class ConnectorManager(BaseComponent):
             StarRocksConnector,
         )
         from dbgpt.datasource.rdbms.conn_vertica import VerticaConnector  # noqa: F401
+        from dbgpt.datasource.rdbms.dialect.oceanbase.ob_dialect import (  # noqa: F401
+            OBDialect,
+        )
 
         from .connect_config_db import ConnectConfigEntity  # noqa: F401
 

dbgpt/datasource/rdbms/conn_oceanbase.py

@@ -0,0 +1,37 @@
+"""Ocean base connect."""
+
+import logging
+
+from dbgpt.datasource.rdbms.base import RDBMSConnector
+
+logger = logging.getLogger(__name__)
+
+
+class OceanBaseConnect(RDBMSConnector):
+    """Connect Oceanbase Database fetch MetaData.
+
+    Args:
+    Usage:
+    """
+
+    db_type: str = "oceanbase"
+    db_dialect: str = "mysql"
+    driver: str = "mysql+ob"
+
+    default_db = ["information_schema", "performance_schema", "sys", "mysql"]
+
+    def get_users(self):
+        """Get_users."""
+        return []
+
+    def get_grants(self):
+        """Get_grants."""
+        return []
+
+    def get_collation(self):
+        """Get collation."""
+        return "UTF-8"
+
+    def get_charset(self):
+        """Get_charset."""
+        return "UTF-8"

dbgpt/datasource/rdbms/dialect/oceanbase/__init__.py

@@ -0,0 +1 @@
+"""Oceanbase connect Module."""

dbgpt/datasource/rdbms/dialect/oceanbase/ob_dialect.py

@@ -0,0 +1,26 @@
+"""OB Dialect support."""
+
+from sqlalchemy.dialects import registry
+from sqlalchemy.dialects.mysql import pymysql
+
+
+class OBDialect(pymysql.MySQLDialect_pymysql):
+    """OBDialect expend."""
+
+    def initialize(self, connection):
+        """Ob dialect initialize."""
+        super(OBDialect, self).initialize(connection)
+        self._server_version_info = (5, 7, 19)
+        self.server_version_info = (5, 7, 19)
+
+    def _server_version_info(self, connection):
+        """Ob set fixed version ending compatibility issue."""
+        return (5, 7, 19)
+
+    def get_isolation_level(self, dbapi_connection):
+        """Ob set fixed version ending compatibility issue."""
+        self.server_version_info = (5, 7, 19)
+        return super(OBDialect, self).get_isolation_level(dbapi_connection)
+
+
+registry.register("mysql.ob", __name__, "OBDialect")

dbgpt/model/adapter/hf_adapter.py

@@ -315,7 +315,6 @@ class QwenAdapter(NewHFChatModelAdapter):
 
 
 class Qwen2Adapter(QwenAdapter):
-
     support_4bit: bool = True
     support_8bit: bool = True
 

dbgpt/model/proxy/llms/zhipu.py

@@ -99,7 +99,6 @@ class ZhipuLLMClient(ProxyLLMClient):
         request: ModelRequest,
         message_converter: Optional[MessageConverter] = None,
     ) -> Iterator[ModelOutput]:
-
         request = self.local_covert_message(request, message_converter)
 
         messages = request.to_common_messages(support_system_role=False)

dbgpt/rag/embedding/embeddings.py

@@ -90,9 +90,9 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings):
             ) from exc
 
         kwargs["client"] = sentence_transformers.SentenceTransformer(
-            kwargs.get("model_name"),
+            kwargs.get("model_name") or DEFAULT_MODEL_NAME,
             cache_folder=kwargs.get("cache_folder"),
-            **kwargs.get("model_kwargs"),
+            **(kwargs.get("model_kwargs") or {}),
         )
         super().__init__(**kwargs)
 

dbgpt/rag/embedding/rerank.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, List, Optional, cast
 
 import aiohttp
+import numpy as np
 import requests
 
 from dbgpt._private.pydantic import EXTRA_FORBID, BaseModel, ConfigDict, Field
@@ -16,7 +17,7 @@ class CrossEncoderRerankEmbeddings(BaseModel, RerankEmbeddings):
 
     client: Any  #: :meta private:
     model_name: str = "BAAI/bge-reranker-base"
-    max_length: Optional[int] = None
+    max_length: int
     """Max length for input sequences. Longer sequences will be truncated. If None, max
         length of the model will be used"""
     """Model name to use."""
@@ -33,9 +34,9 @@ class CrossEncoderRerankEmbeddings(BaseModel, RerankEmbeddings):
             )
 
         kwargs["client"] = CrossEncoder(
-            kwargs.get("model_name"),
+            kwargs.get("model_name", "BAAI/bge-reranker-base"),
             max_length=kwargs.get("max_length"),
-            **kwargs.get("model_kwargs"),
+            **(kwargs.get("model_kwargs") or {}),
         )
         super().__init__(**kwargs)
 
@@ -54,7 +55,9 @@ class CrossEncoderRerankEmbeddings(BaseModel, RerankEmbeddings):
         query_content_pairs = [[query, candidate] for candidate in candidates]
         _model = cast(CrossEncoder, self.client)
         rank_scores = _model.predict(sentences=query_content_pairs)
-        return rank_scores.tolist()
+        if isinstance(rank_scores, np.ndarray):
+            rank_scores = rank_scores.tolist()
+        return rank_scores
 
 
 class OpenAPIRerankEmbeddings(BaseModel, RerankEmbeddings):

dbgpt/rag/retriever/rerank.py

@@ -244,8 +244,8 @@ class CrossEncoderRanker(Ranker):
         contents = [candidate.content for candidate in candidates_with_scores]
         query_content_pairs = [
             [
-                query,
-                content,
+                query if query is not None else "",
+                content if content is not None else "",
             ]
             for content in contents
         ]

dbgpt/storage/schema.py

@@ -16,8 +16,8 @@ class DbInfo:
 class DBType(Enum):
     """Database type enumeration."""
 
-    Mysql = DbInfo("mysql")
-    OCeanBase = DbInfo("oceanbase")
+    MySQL = DbInfo("mysql")
+    OceanBase = DbInfo("oceanbase")
     DuckDb = DbInfo("duckdb", True)
     SQLite = DbInfo("sqlite", True)
     Oracle = DbInfo("oracle")

examples/agents/sql_agent_dialogue_example.py

@@ -30,7 +30,6 @@ initialize_tracer("/tmp/agent_trace.jsonl", create_system_app=True)
 
 
 async def main():
-
     from dbgpt.model.proxy import OpenAILLMClient
 
     llm_client = OpenAILLMClient(model_alias="gpt-3.5-turbo")

web/utils/constants.ts

@@ -45,6 +45,7 @@ export const MODEL_ICON_MAP: Record<ModelType, { label: string; icon: string }>
 
 export const dbMapper: Record<DBType, { label: string; icon: string; desc: string }> = {
   mysql: { label: 'MySQL', icon: '/icons/mysql.png', desc: 'Fast, reliable, scalable open-source relational database management system.' },
+  oceanbase: { label: 'OceanBase', icon: '/icons/oceanbase.png', desc: 'An Ultra-Fast & Cost-Effective Distributed SQL Database.' },
   mssql: { label: 'MSSQL', icon: '/icons/mssql.png', desc: 'Powerful, scalable, secure relational database system by Microsoft.' },
   duckdb: { label: 'DuckDB', icon: '/icons/duckdb.png', desc: 'In-memory analytical database with efficient query processing.' },
   sqlite: { label: 'Sqlite', icon: '/icons/sqlite.png', desc: 'Lightweight embedded relational database with simplicity and portability.' },