feat: Client support chatdata (eosphoros-ai#1343)

dbgpt/app/initialization/serve_initialization.py

@@ -45,11 +45,23 @@ def register_serve_apps(system_app: SystemApp, cfg: Config):
     # Register serve app
     system_app.register(FlowServe)
 
+    # ################################ Rag Serve Register Begin ######################################
+
     from dbgpt.serve.rag.serve import (
         SERVE_CONFIG_KEY_PREFIX as RAG_SERVE_CONFIG_KEY_PREFIX,
     )
     from dbgpt.serve.rag.serve import Serve as RagServe
 
     # Register serve app
     system_app.register(RagServe)
+
+    # ################################ Datasource Serve Register Begin ######################################
+
+    from dbgpt.serve.datasource.serve import (
+        SERVE_CONFIG_KEY_PREFIX as DATASOURCE_SERVE_CONFIG_KEY_PREFIX,
+    )
+    from dbgpt.serve.datasource.serve import Serve as DatasourceServe
+
+    # Register serve app
+    system_app.register(DatasourceServe)
     # ################################ AWEL Flow Serve Register End ########################################

dbgpt/app/openapi/api_v2.py

@@ -127,6 +127,7 @@ async def chat_completions(
         request.chat_mode is None
         or request.chat_mode == ChatMode.CHAT_NORMAL.value
         or request.chat_mode == ChatMode.CHAT_KNOWLEDGE.value
+        or request.chat_mode == ChatMode.CHAT_DATA.value
     ):
         with root_tracer.start_span(
             "get_chat_instance", span_type=SpanType.CHAT, metadata=request.dict()
@@ -146,7 +147,7 @@ async def chat_completions(
             status_code=400,
             detail={
                 "error": {
-                    "message": "chat mode now only support chat_normal, chat_app, chat_flow, chat_knowledge",
+                    "message": "chat mode now only support chat_normal, chat_app, chat_flow, chat_knowledge, chat_data",
                     "type": "invalid_request_error",
                     "param": None,
                     "code": "invalid_chat_mode",
@@ -169,7 +170,8 @@ async def get_chat_instance(dialogue: ChatCompletionRequestBody = Body()) -> Bas
             dialogue.chat_mode, dialogue.user_name, dialogue.sys_code
         )
         dialogue.conv_uid = conv_vo.conv_uid
-
+    if dialogue.chat_mode == "chat_data":
+        dialogue.chat_mode = ChatScene.ChatWithDbExecute.value()
     if not ChatScene.is_valid_mode(dialogue.chat_mode):
         raise StopAsyncIteration(f"Unsupported Chat Mode,{dialogue.chat_mode}!")
 
@@ -201,7 +203,7 @@ async def no_stream_wrapper(
     """
     with root_tracer.start_span("no_stream_generator"):
         response = await chat.nostream_call()
-        msg = response.replace("\ufffd", "")
+        msg = response.replace("\ufffd", "").replace(""", '"')
         choice_data = ChatCompletionResponseChoice(
             index=0,
             message=ChatMessage(role="assistant", content=msg),

dbgpt/client/datasource.py

@@ -0,0 +1,119 @@
+"""this module contains the datasource client functions."""
+from typing import List
+
+from dbgpt.core.schema.api import Result
+
+from .client import Client, ClientException
+from .schema import DatasourceModel
+
+
+async def create_datasource(
+    client: Client, datasource: DatasourceModel
+) -> DatasourceModel:
+    """Create a new datasource.
+
+    Args:
+        client (Client): The dbgpt client.
+        datasource (DatasourceModel): The datasource model.
+    """
+    try:
+        res = await client.get("/datasources", datasource.dict())
+        result: Result = res.json()
+        if result["success"]:
+            return DatasourceModel(**result["data"])
+        else:
+            raise ClientException(status=result["err_code"], reason=result)
+    except Exception as e:
+        raise ClientException(f"Failed to create datasource: {e}")
+
+
+async def update_datasource(
+    client: Client, datasource: DatasourceModel
+) -> DatasourceModel:
+    """Update a datasource.
+
+    Args:
+        client (Client): The dbgpt client.
+        datasource (DatasourceModel): The datasource model.
+    Returns:
+        DatasourceModel: The datasource model.
+    Raises:
+        ClientException: If the request failed.
+    """
+    try:
+        res = await client.put("/datasources", datasource.dict())
+        result: Result = res.json()
+        if result["success"]:
+            return DatasourceModel(**result["data"])
+        else:
+            raise ClientException(status=result["err_code"], reason=result)
+    except Exception as e:
+        raise ClientException(f"Failed to update datasource: {e}")
+
+
+async def delete_datasource(client: Client, datasource_id: str) -> DatasourceModel:
+    """
+    Delete a datasource.
+
+    Args:
+        client (Client): The dbgpt client.
+        datasource_id (str): The datasource id.
+    Returns:
+        DatasourceModel: The datasource model.
+    Raises:
+        ClientException: If the request failed.
+    """
+    try:
+        res = await client.delete("/datasources/" + datasource_id)
+        result: Result = res.json()
+        if result["success"]:
+            return DatasourceModel(**result["data"])
+        else:
+            raise ClientException(status=result["err_code"], reason=result)
+    except Exception as e:
+        raise ClientException(f"Failed to delete datasource: {e}")
+
+
+async def get_datasource(client: Client, datasource_id: str) -> DatasourceModel:
+    """
+    Get a datasource.
+
+    Args:
+        client (Client): The dbgpt client.
+        datasource_id (str): The datasource id.
+    Returns:
+        DatasourceModel: The datasource model.
+    Raises:
+        ClientException: If the request failed.
+    """
+    try:
+        res = await client.get("/datasources/" + datasource_id)
+        result: Result = res.json()
+        if result["success"]:
+            return DatasourceModel(**result["data"])
+        else:
+            raise ClientException(status=result["err_code"], reason=result)
+    except Exception as e:
+        raise ClientException(f"Failed to get datasource: {e}")
+
+
+async def list_datasource(client: Client) -> List[DatasourceModel]:
+    """
+    List datasources.
+
+    Args:
+        client (Client): The dbgpt client.
+    Returns:
+        List[DatasourceModel]: The list of datasource models.
+    Raises:
+        ClientException: If the request failed.
+    """
+    try:
+        res = await client.get("/datasources")
+        result: Result = res.json()
+        if result["success"]:
+            return [DatasourceModel(**datasource) for datasource in result["data"]]
+        else:
+            raise ClientException(status=result["err_code"], reason=result)
+    except Exception as e:
+        raise ClientException(f"Failed to list datasource: {e}")

dbgpt/client/schema.py

@@ -72,6 +72,7 @@ class ChatMode(Enum):
     CHAT_APP = "chat_app"
     CHAT_AWEL_FLOW = "chat_flow"
     CHAT_KNOWLEDGE = "chat_knowledge"
+    CHAT_DATA = "chat_data"
 
 
 class AwelTeamModel(BaseModel):
@@ -278,3 +279,17 @@ class SyncModel(BaseModel):
     """chunk_parameters: chunk parameters
     """
     chunk_parameters: ChunkParameters = Field(None, description="chunk parameters")
+
+
+class DatasourceModel(BaseModel):
+    """Datasource model."""
+
+    id: Optional[int] = Field(None, description="The datasource id")
+    db_type: str = Field(..., description="Database type, e.g. sqlite, mysql, etc.")
+    db_name: str = Field(..., description="Database name.")
+    db_path: str = Field("", description="File path for file-based database.")
+    db_host: str = Field("", description="Database host.")
+    db_port: int = Field(0, description="Database port.")
+    db_user: str = Field("", description="Database user.")
+    db_pwd: str = Field("", description="Database password.")
+    comment: str = Field("", description="Comment for the database.")

dbgpt/datasource/manages/connect_config_db.py

@@ -1,10 +1,14 @@
 """DB Model for connect_config."""
 
 import logging
-from typing import Optional
+from typing import Any, Dict, Optional, Union
 
 from sqlalchemy import Column, Index, Integer, String, Text, UniqueConstraint, text
 
+from dbgpt.serve.datasource.api.schemas import (
+    DatasourceServeRequest,
+    DatasourceServeResponse,
+)
 from dbgpt.storage.metadata import BaseDao, Model
 
 logger = logging.getLogger(__name__)
@@ -218,3 +222,62 @@ def delete_db(self, db_name):
         session.commit()
         session.close()
         return True
+
+    def from_request(
+        self, request: Union[DatasourceServeRequest, Dict[str, Any]]
+    ) -> ConnectConfigEntity:
+        """Convert the request to an entity.
+
+        Args:
+            request (Union[ServeRequest, Dict[str, Any]]): The request
+
+        Returns:
+            T: The entity
+        """
+        request_dict = (
+            request.dict() if isinstance(request, DatasourceServeRequest) else request
+        )
+        entity = ConnectConfigEntity(**request_dict)
+        return entity
+
+    def to_request(self, entity: ConnectConfigEntity) -> DatasourceServeRequest:
+        """Convert the entity to a request.
+
+        Args:
+            entity (T): The entity
+
+        Returns:
+            REQ: The request
+        """
+        return DatasourceServeRequest(
+            id=entity.id,
+            db_type=entity.db_type,
+            db_name=entity.db_name,
+            db_path=entity.db_path,
+            db_host=entity.db_host,
+            db_port=entity.db_port,
+            db_user=entity.db_user,
+            db_pwd=entity.db_pwd,
+            comment=entity.comment,
+        )
+
+    def to_response(self, entity: ConnectConfigEntity) -> DatasourceServeResponse:
+        """Convert the entity to a response.
+
+        Args:
+            entity (T): The entity
+
+        Returns:
+            REQ: The request
+        """
+        return DatasourceServeResponse(
+            id=entity.id,
+            db_type=entity.db_type,
+            db_name=entity.db_name,
+            db_path=entity.db_path,
+            db_host=entity.db_host,
+            db_port=entity.db_port,
+            db_user=entity.db_user,
+            db_pwd=entity.db_pwd,
+            comment=entity.comment,
+        )