Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Client support chatdata #1343

Merged
merged 5 commits into from
Mar 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions dbgpt/app/initialization/serve_initialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 ########################################
8 changes: 5 additions & 3 deletions dbgpt/app/openapi/api_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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",
Expand All @@ -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}!")

Expand Down Expand Up @@ -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),
Expand Down
119 changes: 119 additions & 0 deletions dbgpt/client/datasource.py
Original file line number Diff line number Diff line change
@@ -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}")
15 changes: 15 additions & 0 deletions dbgpt/client/schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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.")
65 changes: 64 additions & 1 deletion dbgpt/datasource/manages/connect_config_db.py
Original file line number Diff line number Diff line change
@@ -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__)
Expand Down Expand Up @@ -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,
)
Empty file.
Loading
Loading