feat: Client support chatdata (#1343)

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.")