mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-07-24 04:36:23 +00:00
feat: Client support chatdata (#1343)
This commit is contained in:
parent
f144fc3d36
commit
dffd235bfb
@ -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 ########################################
|
||||
|
@ -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),
|
||||
|
119
dbgpt/client/datasource.py
Normal file
119
dbgpt/client/datasource.py
Normal file
@ -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}")
|
@ -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.")
|
||||
|
@ -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 @@ class ConnectConfigDao(BaseDao):
|
||||
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,
|
||||
)
|
||||
|
0
dbgpt/serve/datasource/api/__init__.py
Normal file
0
dbgpt/serve/datasource/api/__init__.py
Normal file
193
dbgpt/serve/datasource/api/endpoints.py
Normal file
193
dbgpt/serve/datasource/api/endpoints.py
Normal file
@ -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
|
41
dbgpt/serve/datasource/api/schemas.py
Normal file
41
dbgpt/serve/datasource/api/schemas.py
Normal file
@ -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}"
|
28
dbgpt/serve/datasource/config.py
Normal file
28
dbgpt/serve/datasource/config.py
Normal file
@ -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"},
|
||||
)
|
1
dbgpt/serve/datasource/dependencies.py
Normal file
1
dbgpt/serve/datasource/dependencies.py
Normal file
@ -0,0 +1 @@
|
||||
# Define your dependencies here
|
0
dbgpt/serve/datasource/models/__init__.py
Normal file
0
dbgpt/serve/datasource/models/__init__.py
Normal file
0
dbgpt/serve/datasource/models/models.py
Normal file
0
dbgpt/serve/datasource/models/models.py
Normal file
60
dbgpt/serve/datasource/serve.py
Normal file
60
dbgpt/serve/datasource/serve.py
Normal file
@ -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()
|
0
dbgpt/serve/datasource/service/__init__.py
Normal file
0
dbgpt/serve/datasource/service/__init__.py
Normal file
181
dbgpt/serve/datasource/service/service.py
Normal file
181
dbgpt/serve/datasource/service/service.py
Normal file
@ -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]
|
0
dbgpt/serve/datasource/tests/__init__.py
Normal file
0
dbgpt/serve/datasource/tests/__init__.py
Normal file
@ -166,13 +166,9 @@ Return <a href="#the-app-object">App Object</a> List
|
||||
|
||||
### The App Model
|
||||
________
|
||||
<b>id</b> <font color="gray"> string </font>
|
||||
|
||||
space id
|
||||
________
|
||||
<b>app_code</b> <font color="gray"> string </font>
|
||||
|
||||
app code
|
||||
unique app id
|
||||
________
|
||||
<b>app_name</b> <font color="gray"> string </font>
|
||||
|
||||
|
300
docs/docs/api/datasource.md
Normal file
300
docs/docs/api/datasource.md
Normal file
@ -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
|
||||
|
||||
|
||||
<Tabs
|
||||
defaultValue="python"
|
||||
groupId="chat"
|
||||
values={[
|
||||
{label: 'Curl', value: 'curl'},
|
||||
{label: 'Python', value: 'python'},
|
||||
]
|
||||
}>
|
||||
|
||||
<TabItem value="curl">
|
||||
|
||||
```shell
|
||||
DBGPT_API_KEY=dbgpt
|
||||
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\": \"$DB_NAME\"}"
|
||||
|
||||
```
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="python">
|
||||
|
||||
```python
|
||||
from dbgpt.client import Client
|
||||
|
||||
DBGPT_API_KEY = "dbgpt"
|
||||
DB_NAME="{your_db_name}"
|
||||
|
||||
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
|
||||
)
|
||||
```
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
#### 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<chart-view content=\"{\"type\": \"response_table\", \"sql\": \"SELECT * FROM knowledge_space LIMIT 5\", \"data\": [{\"id\": 5, \"name\": \"frfrw\", \"vector_type\": \"Chroma\", \"desc\": \"eee\", \"owner\": \"eee\", \"context\": null, \"gmt_created\": \"2024-01-02T13:29:52\", \"gmt_modified\": \"2024-01-02T13:29:52\", \"description\": null}, {\"id\": 7, \"name\": \"acc\", \"vector_type\": \"Chroma\", \"desc\": \"dede\", \"owner\": \"dede\", \"context\": null, \"gmt_created\": \"2024-01-02T13:47:01\", \"gmt_modified\": \"2024-01-02T13:47:01\", \"description\": null}, {\"id\": 8, \"name\": \"bcc\", \"vector_type\": \"Chroma\", \"desc\": \"dede\", \"owner\": \"dede\", \"context\": null, \"gmt_created\": \"2024-01-02T14:22:02\", \"gmt_modified\": \"2024-01-02T14:22:02\", \"description\": null}, {\"id\": 9, \"name\": \"dede\", \"vector_type\": \"Chroma\", \"desc\": \"dede\", \"owner\": \"dede\", \"context\": null, \"gmt_created\": \"2024-01-02T14:36:18\", \"gmt_modified\": \"2024-01-02T14:36:18\", \"description\": null}, {\"id\": 10, \"name\": \"qqq\", \"vector_type\": \"Chroma\", \"desc\": \"dede\", \"owner\": \"dede\", \"context\": null, \"gmt_created\": \"2024-01-02T14:40:56\", \"gmt_modified\": \"2024-01-02T14:40:56\", \"description\": null}]}\" />"
|
||||
},
|
||||
"finish_reason": null
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
"completion_tokens": 0
|
||||
}
|
||||
}
|
||||
```
|
||||
### Create Datasource
|
||||
|
||||
```python
|
||||
POST /api/v2/serve/datasources
|
||||
```
|
||||
#### Request body
|
||||
Request <a href="#the-flow-object">Datasource Object</a>
|
||||
|
||||
#### Response body
|
||||
Return <a href="#the-flow-object">Datasource Object</a>
|
||||
|
||||
|
||||
### Update Datasource
|
||||
```python
|
||||
PUT /api/v2/serve/datasources
|
||||
```
|
||||
|
||||
#### Request body
|
||||
Request <a href="#the-flow-object">Datasource Object</a>
|
||||
|
||||
#### Response body
|
||||
Return <a href="#the-flow-object">Datasource Object</a>
|
||||
|
||||
### Delete Datasource
|
||||
|
||||
```python
|
||||
DELETE /api/v2/serve/datasources
|
||||
```
|
||||
|
||||
<Tabs
|
||||
defaultValue="curl_update_datasource"
|
||||
groupId="chat1"
|
||||
values={[
|
||||
{label: 'Curl', value: 'curl_update_datasource'},
|
||||
{label: 'Python', value: 'python_update_datasource'},
|
||||
]
|
||||
}>
|
||||
|
||||
<TabItem value="curl_update_datasource">
|
||||
|
||||
```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" \
|
||||
|
||||
```
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="python_update_datasource">
|
||||
|
||||
|
||||
```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)
|
||||
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
#### Delete Parameters
|
||||
________
|
||||
<b>datasource_id</b> <font color="gray"> string </font> <font color="red"> Required </font>
|
||||
|
||||
datasource id
|
||||
________
|
||||
|
||||
#### Response body
|
||||
Return <a href="#the-flow-object">Datasource Object</a>
|
||||
|
||||
### Get Datasource
|
||||
|
||||
```python
|
||||
GET /api/v2/serve/datasources/{datasource_id}
|
||||
```
|
||||
<Tabs
|
||||
defaultValue="curl_get_datasource"
|
||||
groupId="chat1"
|
||||
values={[
|
||||
{label: 'Curl', value: 'curl_get_datasource'},
|
||||
{label: 'Python', value: 'python_get_datasource'},
|
||||
]
|
||||
}>
|
||||
|
||||
<TabItem value="curl_get_datasource">
|
||||
|
||||
```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"
|
||||
|
||||
```
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="python_get_datasource">
|
||||
|
||||
|
||||
```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)
|
||||
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
#### Query Parameters
|
||||
________
|
||||
<b>datasource_id</b> <font color="gray"> string </font> <font color="red"> Required </font>
|
||||
|
||||
datasource id
|
||||
________
|
||||
|
||||
#### Response body
|
||||
Return <a href="#the-flow-object">Datasource Object</a>
|
||||
|
||||
### List Datasource
|
||||
|
||||
```python
|
||||
GET /api/v2/serve/datasources
|
||||
```
|
||||
|
||||
|
||||
<Tabs
|
||||
defaultValue="curl_list_datasource"
|
||||
groupId="chat1"
|
||||
values={[
|
||||
{label: 'Curl', value: 'curl_list_datasource'},
|
||||
{label: 'Python', value: 'python_list_datasource'},
|
||||
]
|
||||
}>
|
||||
|
||||
<TabItem value="curl_list_datasource">
|
||||
|
||||
```shell
|
||||
DBGPT_API_KEY=dbgpt
|
||||
|
||||
curl -X GET "http://localhost:5000/api/v2/serve/datasources" -H "Authorization: Bearer $DBGPT_API_KEY"
|
||||
|
||||
```
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="python_list_datasource">
|
||||
|
||||
|
||||
```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)
|
||||
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
#### Response body
|
||||
Return <a href="#the-flow-object">Datasource Object</a>
|
||||
|
||||
### The Datasource Object
|
||||
|
||||
________
|
||||
<b>id</b> <font color="gray">string</font>
|
||||
|
||||
The unique id for the datasource.
|
||||
________
|
||||
<b>db_name</b> <font color="gray">string</font>
|
||||
|
||||
The Database name
|
||||
________
|
||||
<b>db_type</b> <font color="gray">string</font>
|
||||
|
||||
Database type, e.g. sqlite, mysql, etc.
|
||||
________
|
||||
<b>db_path</b> <font color="gray">string</font>
|
||||
|
||||
File path for file-based database.
|
||||
________
|
||||
<b>db_host</b> <font color="gray">string</font>
|
||||
|
||||
Database host.
|
||||
________
|
||||
<b>db_port</b> <font color="gray">object</font>
|
||||
|
||||
Database port.
|
||||
________
|
||||
<b>db_user</b> <font color="gray">string</font>
|
||||
|
||||
Database user.
|
||||
________
|
||||
<b>db_pwd</b> <font color="gray">string</font>
|
||||
|
||||
Database password.
|
||||
________
|
||||
<b>comment</b> <font color="gray">string</font>
|
||||
|
||||
Comment for the database.
|
||||
________
|
@ -1,6 +1,6 @@
|
||||
# Flow
|
||||
|
||||
Get started with the App API
|
||||
Get started with the Flow API
|
||||
|
||||
# Chat Flow
|
||||
|
||||
@ -76,8 +76,9 @@ Return <a href="#the-flow-object">Flow Object</a>
|
||||
|
||||
|
||||
### Update Flow
|
||||
|
||||
```python
|
||||
PUT /api/v2/serve/awel/flows
|
||||
```
|
||||
|
||||
#### Request body
|
||||
Request <a href="#the-flow-object">Flow Object</a>
|
||||
@ -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}"
|
||||
|
@ -389,6 +389,9 @@ const sidebars = {
|
||||
{
|
||||
type: 'doc',
|
||||
id: 'api/knowledge'
|
||||
},{
|
||||
type: 'doc',
|
||||
id: 'api/datasource'
|
||||
},
|
||||
],
|
||||
link: {
|
||||
|
62
examples/client/datasource_crud_example.py
Normal file
62
examples/client/datasource_crud_example.py
Normal file
@ -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())
|
Loading…
Reference in New Issue
Block a user