From 44d9f7c7da3a334f72280121883cb478467950c4 Mon Sep 17 00:00:00 2001
From: aries_ckt <916701291@qq.com>
Date: Wed, 27 Mar 2024 16:04:30 +0800
Subject: [PATCH 1/4] feat:add client datasource
---
.../initialization/serve_initialization.py | 12 +
dbgpt/app/openapi/api_v2.py | 8 +-
dbgpt/client/datasource.py | 119 +++++++
dbgpt/client/schema.py | 15 +
dbgpt/datasource/manages/connect_config_db.py | 65 +++-
dbgpt/serve/datasource/api/__init__.py | 0
dbgpt/serve/datasource/api/endpoints.py | 193 +++++++++++
dbgpt/serve/datasource/api/schemas.py | 41 +++
dbgpt/serve/datasource/config.py | 28 ++
dbgpt/serve/datasource/dependencies.py | 1 +
dbgpt/serve/datasource/models/__init__.py | 0
dbgpt/serve/datasource/models/models.py | 0
dbgpt/serve/datasource/serve.py | 60 ++++
dbgpt/serve/datasource/service/__init__.py | 0
dbgpt/serve/datasource/service/service.py | 181 +++++++++++
dbgpt/serve/datasource/tests/__init__.py | 0
docs/docs/api/app.md | 6 +-
docs/docs/api/datasource.md | 300 ++++++++++++++++++
docs/docs/api/flow.md | 7 +-
docs/sidebars.js | 3 +
examples/client/datasource_crud_example.py | 62 ++++
examples/client/flow_crud_example.py | 2 +
22 files changed, 1091 insertions(+), 12 deletions(-)
create mode 100644 dbgpt/client/datasource.py
create mode 100644 dbgpt/serve/datasource/api/__init__.py
create mode 100644 dbgpt/serve/datasource/api/endpoints.py
create mode 100644 dbgpt/serve/datasource/api/schemas.py
create mode 100644 dbgpt/serve/datasource/config.py
create mode 100644 dbgpt/serve/datasource/dependencies.py
create mode 100644 dbgpt/serve/datasource/models/__init__.py
create mode 100644 dbgpt/serve/datasource/models/models.py
create mode 100644 dbgpt/serve/datasource/serve.py
create mode 100644 dbgpt/serve/datasource/service/__init__.py
create mode 100644 dbgpt/serve/datasource/service/service.py
create mode 100644 dbgpt/serve/datasource/tests/__init__.py
create mode 100644 docs/docs/api/datasource.md
create mode 100644 examples/client/datasource_crud_example.py
diff --git a/dbgpt/app/initialization/serve_initialization.py b/dbgpt/app/initialization/serve_initialization.py
index 4e757fca5..106da8fc9 100644
--- a/dbgpt/app/initialization/serve_initialization.py
+++ b/dbgpt/app/initialization/serve_initialization.py
@@ -45,6 +45,8 @@ 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,
)
@@ -52,4 +54,14 @@ def register_serve_apps(system_app: SystemApp, cfg: Config):
# 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 ########################################
diff --git a/dbgpt/app/openapi/api_v2.py b/dbgpt/app/openapi/api_v2.py
index 14cf2f554..a43a895ba 100644
--- a/dbgpt/app/openapi/api_v2.py
+++ b/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),
diff --git a/dbgpt/client/datasource.py b/dbgpt/client/datasource.py
new file mode 100644
index 000000000..47244b47e
--- /dev/null
+++ b/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}")
diff --git a/dbgpt/client/schema.py b/dbgpt/client/schema.py
index b4e8ad397..7043dc427 100644
--- a/dbgpt/client/schema.py
+++ b/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.")
diff --git a/dbgpt/datasource/manages/connect_config_db.py b/dbgpt/datasource/manages/connect_config_db.py
index 866bc1065..ac80ae1a7 100644
--- a/dbgpt/datasource/manages/connect_config_db.py
+++ b/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,
+ )
diff --git a/dbgpt/serve/datasource/api/__init__.py b/dbgpt/serve/datasource/api/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/dbgpt/serve/datasource/api/endpoints.py b/dbgpt/serve/datasource/api/endpoints.py
new file mode 100644
index 000000000..4401b32b4
--- /dev/null
+++ b/dbgpt/serve/datasource/api/endpoints.py
@@ -0,0 +1,193 @@
+from functools import cache
+from typing import List, Optional
+
+from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBearer
+
+from dbgpt.component import SystemApp
+from dbgpt.serve.core import Result
+from dbgpt.serve.datasource.api.schemas import (
+ DatasourceServeRequest,
+ DatasourceServeResponse,
+)
+from dbgpt.serve.datasource.config import SERVE_SERVICE_COMPONENT_NAME
+from dbgpt.serve.datasource.service.service import Service
+from dbgpt.util import PaginationResult
+
+router = APIRouter()
+
+# Add your API endpoints here
+
+global_system_app: Optional[SystemApp] = None
+
+
+def get_service() -> Service:
+ """Get the service instance"""
+ return global_system_app.get_component(SERVE_SERVICE_COMPONENT_NAME, Service)
+
+
+get_bearer_token = HTTPBearer(auto_error=False)
+
+
+@cache
+def _parse_api_keys(api_keys: str) -> List[str]:
+ """Parse the string api keys to a list
+
+ Args:
+ api_keys (str): The string api keys
+
+ Returns:
+ List[str]: The list of api keys
+ """
+ if not api_keys:
+ return []
+ return [key.strip() for key in api_keys.split(",")]
+
+
+async def check_api_key(
+ auth: Optional[HTTPAuthorizationCredentials] = Depends(get_bearer_token),
+ service: Service = Depends(get_service),
+) -> Optional[str]:
+ """Check the api key
+
+ If the api key is not set, allow all.
+
+ Your can pass the token in you request header like this:
+
+ .. code-block:: python
+
+ import requests
+
+ client_api_key = "your_api_key"
+ headers = {"Authorization": "Bearer " + client_api_key}
+ res = requests.get("http://test/hello", headers=headers)
+ assert res.status_code == 200
+
+ """
+ if service.config.api_keys:
+ api_keys = _parse_api_keys(service.config.api_keys)
+ if auth is None or (token := auth.credentials) not in api_keys:
+ raise HTTPException(
+ status_code=401,
+ detail={
+ "error": {
+ "message": "",
+ "type": "invalid_request_error",
+ "param": None,
+ "code": "invalid_api_key",
+ }
+ },
+ )
+ return token
+ else:
+ # api_keys not set; allow all
+ return None
+
+
+@router.get("/health", dependencies=[Depends(check_api_key)])
+async def health():
+ """Health check endpoint"""
+ return {"status": "ok"}
+
+
+@router.get("/test_auth", dependencies=[Depends(check_api_key)])
+async def test_auth():
+ """Test auth endpoint"""
+ return {"status": "ok"}
+
+
+@router.post("/datasources", dependencies=[Depends(check_api_key)])
+async def create(
+ request: DatasourceServeRequest, service: Service = Depends(get_service)
+) -> Result:
+ """Create a new Space entity
+
+ Args:
+ request (DatasourceServeRequest): The request
+ service (Service): The service
+ Returns:
+ ServerResponse: The response
+ """
+ return Result.succ(service.create(request))
+
+
+@router.put("/datasources", dependencies=[Depends(check_api_key)])
+async def update(
+ request: DatasourceServeRequest, service: Service = Depends(get_service)
+) -> Result:
+ """Update a Space entity
+
+ Args:
+ request (DatasourceServeRequest): The request
+ service (Service): The service
+ Returns:
+ ServerResponse: The response
+ """
+ return Result.succ(service.update(request))
+
+
+@router.delete(
+ "/datasources/{datasource_id}",
+ response_model=Result[None],
+ dependencies=[Depends(check_api_key)],
+)
+async def delete(
+ datasource_id: str, service: Service = Depends(get_service)
+) -> Result[None]:
+ """Delete a Space entity
+
+ Args:
+ request (DatasourceServeRequest): The request
+ service (Service): The service
+ Returns:
+ ServerResponse: The response
+ """
+ return Result.succ(service.delete(datasource_id))
+
+
+@router.get(
+ "/datasources/{datasource_id}",
+ dependencies=[Depends(check_api_key)],
+ response_model=Result[List],
+)
+async def query(
+ datasource_id: str, service: Service = Depends(get_service)
+) -> Result[List[DatasourceServeResponse]]:
+ """Query Space entities
+
+ Args:
+ request (DatasourceServeRequest): The request
+ service (Service): The service
+ Returns:
+ List[ServeResponse]: The response
+ """
+ return Result.succ(service.get(datasource_id))
+
+
+@router.get(
+ "/datasources",
+ dependencies=[Depends(check_api_key)],
+ response_model=Result[PaginationResult[DatasourceServeResponse]],
+)
+async def query_page(
+ page: int = Query(default=1, description="current page"),
+ page_size: int = Query(default=20, description="page size"),
+ service: Service = Depends(get_service),
+) -> Result[PaginationResult[DatasourceServeResponse]]:
+ """Query Space entities
+
+ Args:
+ page (int): The page number
+ page_size (int): The page size
+ service (Service): The service
+ Returns:
+ ServerResponse: The response
+ """
+ return Result.succ(service.list())
+
+
+def init_endpoints(system_app: SystemApp) -> None:
+ """Initialize the endpoints"""
+ global global_system_app
+ system_app.register(Service)
+ global_system_app = system_app
diff --git a/dbgpt/serve/datasource/api/schemas.py b/dbgpt/serve/datasource/api/schemas.py
new file mode 100644
index 000000000..907ed4f6e
--- /dev/null
+++ b/dbgpt/serve/datasource/api/schemas.py
@@ -0,0 +1,41 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from ..config import SERVE_APP_NAME_HUMP
+
+
+class DatasourceServeRequest(BaseModel):
+ """name: knowledge space name"""
+
+ """vector_type: vector type"""
+ 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.")
+
+
+class DatasourceServeResponse(BaseModel):
+ """Flow response model"""
+
+ """name: knowledge space name"""
+
+ """vector_type: vector type"""
+ id: 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.")
+
+ # TODO define your own fields here
+ class Config:
+ title = f"ServerResponse for {SERVE_APP_NAME_HUMP}"
diff --git a/dbgpt/serve/datasource/config.py b/dbgpt/serve/datasource/config.py
new file mode 100644
index 000000000..d53ed2682
--- /dev/null
+++ b/dbgpt/serve/datasource/config.py
@@ -0,0 +1,28 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+from dbgpt.serve.core import BaseServeConfig
+
+APP_NAME = "datasource"
+SERVE_APP_NAME = "dbgpt_datasource"
+SERVE_APP_NAME_HUMP = "dbgpt_datasource"
+SERVE_CONFIG_KEY_PREFIX = "dbgpt_datasource"
+SERVE_SERVICE_COMPONENT_NAME = f"{SERVE_APP_NAME}_service"
+
+
+@dataclass
+class ServeConfig(BaseServeConfig):
+ """Parameters for the serve command"""
+
+ api_keys: Optional[str] = field(
+ default=None, metadata={"help": "API keys for the endpoint, if None, allow all"}
+ )
+
+ default_user: Optional[str] = field(
+ default=None,
+ metadata={"help": "Default user name for prompt"},
+ )
+ default_sys_code: Optional[str] = field(
+ default=None,
+ metadata={"help": "Default system code for prompt"},
+ )
diff --git a/dbgpt/serve/datasource/dependencies.py b/dbgpt/serve/datasource/dependencies.py
new file mode 100644
index 000000000..8598ecd97
--- /dev/null
+++ b/dbgpt/serve/datasource/dependencies.py
@@ -0,0 +1 @@
+# Define your dependencies here
diff --git a/dbgpt/serve/datasource/models/__init__.py b/dbgpt/serve/datasource/models/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/dbgpt/serve/datasource/models/models.py b/dbgpt/serve/datasource/models/models.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/dbgpt/serve/datasource/serve.py b/dbgpt/serve/datasource/serve.py
new file mode 100644
index 000000000..531080e9d
--- /dev/null
+++ b/dbgpt/serve/datasource/serve.py
@@ -0,0 +1,60 @@
+import logging
+from typing import List, Optional, Union
+
+from sqlalchemy import URL
+
+from dbgpt.component import SystemApp
+from dbgpt.serve.core import BaseServe
+from dbgpt.storage.metadata import DatabaseManager
+
+from .api.endpoints import init_endpoints, router
+from .config import (
+ APP_NAME,
+ SERVE_APP_NAME,
+ SERVE_APP_NAME_HUMP,
+ SERVE_CONFIG_KEY_PREFIX,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Serve(BaseServe):
+ """Serve component for DB-GPT"""
+
+ name = SERVE_APP_NAME
+
+ def __init__(
+ self,
+ system_app: SystemApp,
+ api_prefix: Optional[str] = f"/api/v2/serve",
+ api_tags: Optional[List[str]] = None,
+ db_url_or_db: Union[str, URL, DatabaseManager] = None,
+ try_create_tables: Optional[bool] = False,
+ ):
+ if api_tags is None:
+ api_tags = [SERVE_APP_NAME_HUMP]
+ super().__init__(
+ system_app, api_prefix, api_tags, db_url_or_db, try_create_tables
+ )
+ self._db_manager: Optional[DatabaseManager] = None
+
+ def init_app(self, system_app: SystemApp):
+ if self._app_has_initiated:
+ return
+ self._system_app = system_app
+ self._system_app.app.include_router(
+ router, prefix=self._api_prefix, tags=self._api_tags
+ )
+ init_endpoints(self._system_app)
+ self._app_has_initiated = True
+
+ def on_init(self):
+ """Called when init the application.
+
+ You can do some initialization here. You can't get other components here because they may be not initialized yet
+ """
+
+ def before_start(self):
+ """Called before the start of the application."""
+ # TODO: Your code here
+ self._db_manager = self.create_or_get_db_manager()
diff --git a/dbgpt/serve/datasource/service/__init__.py b/dbgpt/serve/datasource/service/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/dbgpt/serve/datasource/service/service.py b/dbgpt/serve/datasource/service/service.py
new file mode 100644
index 000000000..c409c0bb3
--- /dev/null
+++ b/dbgpt/serve/datasource/service/service.py
@@ -0,0 +1,181 @@
+import logging
+from typing import List, Optional
+
+from fastapi import HTTPException
+
+from dbgpt._private.config import Config
+from dbgpt.component import ComponentType, SystemApp
+from dbgpt.core.awel.dag.dag_manager import DAGManager
+from dbgpt.datasource.db_conn_info import DBConfig
+from dbgpt.datasource.manages.connect_config_db import (
+ ConnectConfigDao,
+ ConnectConfigEntity,
+)
+from dbgpt.serve.core import BaseService
+from dbgpt.storage.metadata import BaseDao
+from dbgpt.storage.schema import DBType
+from dbgpt.storage.vector_store.base import VectorStoreConfig
+from dbgpt.storage.vector_store.connector import VectorStoreConnector
+from dbgpt.util.executor_utils import ExecutorFactory
+
+from ..api.schemas import DatasourceServeRequest, DatasourceServeResponse
+from ..config import SERVE_CONFIG_KEY_PREFIX, SERVE_SERVICE_COMPONENT_NAME, ServeConfig
+
+logger = logging.getLogger(__name__)
+CFG = Config()
+
+
+class Service(
+ BaseService[ConnectConfigEntity, DatasourceServeRequest, DatasourceServeResponse]
+):
+ """The service class for Flow"""
+
+ name = SERVE_SERVICE_COMPONENT_NAME
+
+ def __init__(
+ self,
+ system_app: SystemApp,
+ dao: Optional[ConnectConfigDao] = None,
+ ):
+ self._system_app = None
+ self._dao: ConnectConfigDao = dao
+ self._dag_manager: Optional[DAGManager] = None
+ self._db_summary_client = None
+ self._vector_connector = None
+
+ super().__init__(system_app)
+
+ def init_app(self, system_app: SystemApp) -> None:
+ """Initialize the service
+
+ Args:
+ system_app (SystemApp): The system app
+ """
+ self._serve_config = ServeConfig.from_app_config(
+ system_app.config, SERVE_CONFIG_KEY_PREFIX
+ )
+ self._dao = self._dao or ConnectConfigDao()
+ self._system_app = system_app
+
+ def before_start(self):
+ """Execute before the application starts"""
+ from dbgpt.rag.summary.db_summary_client import DBSummaryClient
+
+ self._db_summary_client = DBSummaryClient(self._system_app)
+
+ def after_start(self):
+ """Execute after the application starts"""
+
+ @property
+ def dao(
+ self,
+ ) -> BaseDao[ConnectConfigEntity, DatasourceServeRequest, DatasourceServeResponse]:
+ """Returns the internal DAO."""
+ return self._dao
+
+ @property
+ def config(self) -> ServeConfig:
+ """Returns the internal ServeConfig."""
+ return self._serve_config
+
+ def create(self, request: DatasourceServeRequest) -> DatasourceServeResponse:
+ """Create a new Datasource entity
+
+ Args:
+ request (DatasourceServeRequest): The request
+
+ Returns:
+ DatasourceServeResponse: The response
+ """
+ datasource = self._dao.get_by_names(request.db_name)
+ if datasource:
+ raise HTTPException(
+ status_code=400,
+ detail=f"datasource name:{request.db_name} already exists",
+ )
+ try:
+ db_type = DBType.of_db_type(request.db_type)
+ if not db_type:
+ raise HTTPException(
+ status_code=400, detail=f"Unsupported Db Type, {request.db_type}"
+ )
+ res = self._dao.create(request)
+
+ # async embedding
+ executor = self._system_app.get_component(
+ ComponentType.EXECUTOR_DEFAULT, ExecutorFactory
+ ).create() # type: ignore
+ executor.submit(
+ self._db_summary_client.db_summary_embedding,
+ request.db_name,
+ request.db_type,
+ )
+ except Exception as e:
+ raise ValueError("Add db connect info error!" + str(e))
+ return res
+
+ def update(self, request: DatasourceServeRequest) -> DatasourceServeResponse:
+ """Create a new Datasource entity
+
+ Args:
+ request (DatasourceServeRequest): The request
+
+ Returns:
+ DatasourceServeResponse: The response
+ """
+ datasources = self._dao.get_by_names(request.db_name)
+ if datasources is None:
+ raise HTTPException(
+ status_code=400,
+ detail=f"there is no datasource name:{request.db_name} exists",
+ )
+ db_config = DBConfig(**request.dict())
+ if CFG.local_db_manager.edit_db(db_config):
+ return DatasourceServeResponse(**db_config.dict())
+ else:
+ raise HTTPException(
+ status_code=400,
+ detail=f"update datasource name:{request.db_name} failed",
+ )
+
+ def get(self, datasource_id: str) -> Optional[DatasourceServeResponse]:
+ """Get a Flow entity
+
+ Args:
+ request (DatasourceServeRequest): The request
+
+ Returns:
+ DatasourceServeResponse: The response
+ """
+ return self._dao.get_one({"id": datasource_id})
+
+ def delete(self, datasource_id: str) -> Optional[DatasourceServeResponse]:
+ """Delete a Flow entity
+
+ Args:
+ datasource_id (str): The datasource_id
+
+ Returns:
+ DatasourceServeResponse: The data after deletion
+ """
+ db_config = self._dao.get_one({"id": datasource_id})
+ vector_name = db_config.db_name + "_profile"
+ vector_store_config = VectorStoreConfig(name=vector_name)
+ self._vector_connector = VectorStoreConnector(
+ vector_store_type=CFG.VECTOR_STORE_TYPE,
+ vector_store_config=vector_store_config,
+ )
+ self._vector_connector.delete_vector_name(vector_name)
+ if db_config:
+ self._dao.delete({"id": datasource_id})
+ return db_config
+
+ def list(self) -> List[DatasourceServeResponse]:
+ """List the Flow entities.
+
+ Returns:
+ List[DatasourceServeResponse]: The list of responses
+ """
+
+ db_list = CFG.local_db_manager.get_db_list()
+ return [DatasourceServeResponse(**db) for db in db_list]
diff --git a/dbgpt/serve/datasource/tests/__init__.py b/dbgpt/serve/datasource/tests/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/docs/docs/api/app.md b/docs/docs/api/app.md
index f2cbec390..190af70b0 100644
--- a/docs/docs/api/app.md
+++ b/docs/docs/api/app.md
@@ -166,13 +166,9 @@ Return App Object List
### The App Model
________
-id string
-
-space id
-________
app_code string
-app code
+unique app id
________
app_name string
diff --git a/docs/docs/api/datasource.md b/docs/docs/api/datasource.md
new file mode 100644
index 000000000..0dcd633a5
--- /dev/null
+++ b/docs/docs/api/datasource.md
@@ -0,0 +1,300 @@
+# Datasource
+
+Get started with the Datasource API
+
+# Chat Datasource
+
+```python
+POST /api/v2/chat/completions
+```
+### Examples
+
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+### Chat Datasource
+
+
+
+
+
+
+```shell
+DBGPT_API_KEY=dbgpt
+DATASOURCE_ID={YOUR_DATASOURCE_ID}
+
+curl -X POST "http://localhost:5000/api/v2/chat/completions" \
+ -H "Authorization: Bearer $DBGPT_API_KEY" \
+ -H "accept: application/json" \
+ -H "Content-Type: application/json" \
+ -d "{\"messages\":\"show space datas limit 5\",\"model\":\"chatgpt_proxyllm\", \"chat_mode\": \"chat_datasource\", \"chat_param\": \"$DATASOURCE_ID\"}"
+
+```
+
+
+
+
+```python
+from dbgpt.client import Client
+
+DBGPT_API_KEY = "dbgpt"
+DB_NAME="dbgpt"
+
+client = Client(api_key=DBGPT_API_KEY)
+res = client.chat(
+ messages="show space datas limit 5",
+ model="chatgpt_proxyllm",
+ chat_mode="chat_data",
+ chat_param=DB_NAME
+)
+```
+
+
+
+#### Chat Completion Response
+```json
+{
+ "id": "2bb80fdd-e47e-4083-8bc9-7ca66ee0931b",
+ "object": "chat.completion",
+ "created": 1711509733,
+ "model": "chatgpt_proxyllm",
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "The user wants to display information about knowledge spaces with a limit of 5 results.\\n"
+ },
+ "finish_reason": null
+ }
+ ],
+ "usage": {
+ "prompt_tokens": 0,
+ "total_tokens": 0,
+ "completion_tokens": 0
+ }
+}
+```
+### Create Datasource
+
+```python
+POST /api/v2/serve/datasources
+```
+#### Request body
+Request Datasource Object
+
+#### Response body
+Return Datasource Object
+
+
+### Update Datasource
+```python
+PUT /api/v2/serve/datasources
+```
+
+#### Request body
+Request Datasource Object
+
+#### Response body
+Return Datasource Object
+
+### Delete Datasource
+
+```python
+DELETE /api/v2/serve/datasources
+```
+
+
+
+
+
+```shell
+DBGPT_API_KEY=dbgpt
+DATASOURCE_ID={YOUR_DATASOURCE_ID}
+
+ curl -X DELETE "http://localhost:5000/api/v2/serve/datasources/$DATASOURCE_ID" \
+ -H "Authorization: Bearer $DBGPT_API_KEY" \
+
+```
+
+
+
+
+
+```python
+from dbgpt.client import Client
+from dbgpt.client.datasource import delete_datasource
+
+DBGPT_API_KEY = "dbgpt"
+datasource_id = "{your_datasource_id}"
+
+client = Client(api_key=DBGPT_API_KEY)
+res = await delete_datasource(client=client, datasource_id=datasource_id)
+
+```
+
+
+
+
+#### Delete Parameters
+________
+datasource_id string Required
+
+datasource id
+________
+
+#### Response body
+Return Datasource Object
+
+### Get Datasource
+
+```python
+GET /api/v2/serve/datasources/{datasource_id}
+```
+
+
+
+
+```shell
+DBGPT_API_KEY=dbgpt
+DATASOURCE_ID={YOUR_DATASOURCE_ID}
+
+curl -X GET "http://localhost:5000/api/v2/serve/datasources/$DATASOURCE_ID" -H "Authorization: Bearer $DBGPT_API_KEY"
+
+```
+
+
+
+
+
+```python
+from dbgpt.client import Client
+from dbgpt.client.datasource import get_datasource
+
+DBGPT_API_KEY = "dbgpt"
+datasource_id = "{your_datasource_id}"
+
+client = Client(api_key=DBGPT_API_KEY)
+res = await get_datasource(client=client, datasource_id=datasource_id)
+
+```
+
+
+
+
+#### Query Parameters
+________
+datasource_id string Required
+
+datasource id
+________
+
+#### Response body
+Return Datasource Object
+
+### List Datasource
+
+```python
+GET /api/v2/serve/datasources
+```
+
+
+
+
+
+
+```shell
+DBGPT_API_KEY=dbgpt
+
+curl -X GET "http://localhost:5000/api/v2/serve/datasources" -H "Authorization: Bearer $DBGPT_API_KEY"
+
+```
+
+
+
+
+
+```python
+from dbgpt.client import Client
+from dbgpt.client.datasource import list_datasource
+
+DBGPT_API_KEY = "dbgpt"
+
+client = Client(api_key=DBGPT_API_KEY)
+res = await list_datasource(client=client)
+
+```
+
+
+
+
+#### Response body
+Return Datasource Object
+
+### The Datasource Object
+
+________
+id string
+
+The unique id for the datasource.
+________
+db_name string
+
+The Database name
+________
+db_type string
+
+Database type, e.g. sqlite, mysql, etc.
+________
+db_path string
+
+File path for file-based database.
+________
+db_host string
+
+Database host.
+________
+db_port object
+
+Database port.
+________
+db_user string
+
+Database user.
+________
+db_pwd string
+
+Database password.
+________
+comment string
+
+Comment for the database.
+________
diff --git a/docs/docs/api/flow.md b/docs/docs/api/flow.md
index 88be326d2..f73565a52 100644
--- a/docs/docs/api/flow.md
+++ b/docs/docs/api/flow.md
@@ -1,6 +1,6 @@
# Flow
-Get started with the App API
+Get started with the Flow API
# Chat Flow
@@ -76,8 +76,9 @@ Return Flow Object
### Update Flow
-
+```python
PUT /api/v2/serve/awel/flows
+```
#### Request body
Request Flow Object
@@ -170,7 +171,7 @@ curl -X GET "http://localhost:5000/api/v2/serve/awel/flows/$FLOW_ID" -H "Authori
```python
from dbgpt.client import Client
-from dbgpt.client.knowledge import get_flow
+from dbgpt.client.flow import get_flow
DBGPT_API_KEY = "dbgpt"
flow_id = "{your_flow_id}"
diff --git a/docs/sidebars.js b/docs/sidebars.js
index 00b094e2e..6735f06f2 100755
--- a/docs/sidebars.js
+++ b/docs/sidebars.js
@@ -385,6 +385,9 @@ const sidebars = {
{
type: 'doc',
id: 'api/knowledge'
+ },{
+ type: 'doc',
+ id: 'api/datasource'
},
],
link: {
diff --git a/examples/client/datasource_crud_example.py b/examples/client/datasource_crud_example.py
new file mode 100644
index 000000000..6ea1e7244
--- /dev/null
+++ b/examples/client/datasource_crud_example.py
@@ -0,0 +1,62 @@
+"""Client: Simple Flow CRUD example
+
+This example demonstrates how to use the dbgpt client to create, get, update, and
+delete datasource.
+
+Example:
+ .. code-block:: python
+
+ DBGPT_API_KEY = "dbgpt"
+ client = Client(api_key=DBGPT_API_KEY)
+ # 1. Create a flow
+ res = await create_datasource(
+ client,
+ DatasourceModel(
+ db_name="dbgpt",
+ desc="for client datasource",
+ db_type="mysql",
+ db_type="mysql",
+ db_host="127.0.0.1",
+ db_user="root",
+ db_pwd="xxxx",
+ db_port=3306,
+ ),
+ )
+ # 2. Update a flow
+ res = await update_datasource(
+ client,
+ DatasourceModel(
+ db_name="dbgpt",
+ desc="for client datasource",
+ db_type="mysql",
+ db_type="mysql",
+ db_host="127.0.0.1",
+ db_user="root",
+ db_pwd="xxxx",
+ db_port=3306,
+ ),
+ )
+ # 3. Delete a flow
+ res = await delete_datasource(client, datasource_id="10")
+ # 4. Get a flow
+ res = await get_datasource(client, datasource_id="10")
+ # 5. List all datasource
+ res = await list_datasource(client)
+
+"""
+import asyncio
+
+from dbgpt.client import Client
+from dbgpt.client.datasource import list_datasource
+
+
+async def main():
+ # initialize client
+ DBGPT_API_KEY = "dbgpt"
+ client = Client(api_key=DBGPT_API_KEY)
+ res = await list_datasource(client)
+ print(res)
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
diff --git a/examples/client/flow_crud_example.py b/examples/client/flow_crud_example.py
index 7b83a00a0..fb02096f9 100644
--- a/examples/client/flow_crud_example.py
+++ b/examples/client/flow_crud_example.py
@@ -37,7 +37,9 @@ async def main():
DBGPT_API_KEY = "dbgpt"
client = Client(api_key=DBGPT_API_KEY)
res = await list_flow(client)
+ res = await list_flow(client)
print(res)
+ await client.aclose()
if __name__ == "__main__":
From ed8c9d4a60fa50b6a8abc93ad1e9f87131241031 Mon Sep 17 00:00:00 2001
From: aries_ckt <916701291@qq.com>
Date: Wed, 27 Mar 2024 16:35:51 +0800
Subject: [PATCH 2/4] feat:update flow example
---
examples/client/flow_crud_example.py | 2 --
1 file changed, 2 deletions(-)
diff --git a/examples/client/flow_crud_example.py b/examples/client/flow_crud_example.py
index fb02096f9..7b83a00a0 100644
--- a/examples/client/flow_crud_example.py
+++ b/examples/client/flow_crud_example.py
@@ -37,9 +37,7 @@ async def main():
DBGPT_API_KEY = "dbgpt"
client = Client(api_key=DBGPT_API_KEY)
res = await list_flow(client)
- res = await list_flow(client)
print(res)
- await client.aclose()
if __name__ == "__main__":
From d24cc4aded53d8ce3aa3764daa1eca904700a2c8 Mon Sep 17 00:00:00 2001
From: aries_ckt <916701291@qq.com>
Date: Wed, 27 Mar 2024 16:35:51 +0800
Subject: [PATCH 3/4] feat:update flow example
---
examples/client/flow_crud_example.py | 2 --
1 file changed, 2 deletions(-)
diff --git a/examples/client/flow_crud_example.py b/examples/client/flow_crud_example.py
index fb02096f9..7b83a00a0 100644
--- a/examples/client/flow_crud_example.py
+++ b/examples/client/flow_crud_example.py
@@ -37,9 +37,7 @@ async def main():
DBGPT_API_KEY = "dbgpt"
client = Client(api_key=DBGPT_API_KEY)
res = await list_flow(client)
- res = await list_flow(client)
print(res)
- await client.aclose()
if __name__ == "__main__":
From 8cb1f6717a992859f213de2e8ce5e408f5602c2e Mon Sep 17 00:00:00 2001
From: aries_ckt <916701291@qq.com>
Date: Wed, 27 Mar 2024 17:22:47 +0800
Subject: [PATCH 4/4] feat:update datasource docs
---
docs/docs/api/datasource.md | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/docs/docs/api/datasource.md b/docs/docs/api/datasource.md
index 0dcd633a5..1e133176c 100644
--- a/docs/docs/api/datasource.md
+++ b/docs/docs/api/datasource.md
@@ -28,13 +28,13 @@ import TabItem from '@theme/TabItem';
```shell
DBGPT_API_KEY=dbgpt
-DATASOURCE_ID={YOUR_DATASOURCE_ID}
+DB_NAME="{your_db_name}"
curl -X POST "http://localhost:5000/api/v2/chat/completions" \
-H "Authorization: Bearer $DBGPT_API_KEY" \
-H "accept: application/json" \
-H "Content-Type: application/json" \
- -d "{\"messages\":\"show space datas limit 5\",\"model\":\"chatgpt_proxyllm\", \"chat_mode\": \"chat_datasource\", \"chat_param\": \"$DATASOURCE_ID\"}"
+ -d "{\"messages\":\"show space datas limit 5\",\"model\":\"chatgpt_proxyllm\", \"chat_mode\": \"chat_datasource\", \"chat_param\": \"$DB_NAME\"}"
```
@@ -45,7 +45,7 @@ curl -X POST "http://localhost:5000/api/v2/chat/completions" \
from dbgpt.client import Client
DBGPT_API_KEY = "dbgpt"
-DB_NAME="dbgpt"
+DB_NAME="{your_db_name}"
client = Client(api_key=DBGPT_API_KEY)
res = client.chat(