github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-13 05:01:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix:client mypy error

Browse Source
This commit is contained in:
aries_ckt
2024-03-20 16:22:38 +08:00
parent 7bc5c59a89
commit f43abf3155
19 changed files with 1814 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
dbgpt/app/openapi/api_v2.py
View File

@@ -72,7 +72,6 @@ async def check_api_key(
@router.post("/v2/chat/completions", dependencies=[Depends(check_api_key)])
async def chat_completions(
request: ChatCompletionRequestBody = Body(),
flow_service: FlowService = Depends(get_chat_flow),
):
"""Chat V2 completions
Args:
@@ -121,7 +120,9 @@ async def chat_completions(
media_type="text/event-stream",
)
elif (
request.chat_mode is None or request.chat_mode == ChatMode.CHAT_KNOWLEDGE.value
request.chat_mode is None
or request.chat_mode == ChatMode.CHAT_NORMAL.value
or request.chat_mode == ChatMode.CHAT_KNOWLEDGE.value
):
with root_tracer.start_span(
"get_chat_instance", span_type=SpanType.CHAT, metadata=request.dict()

38
dbgpt/client/app.py
View File

@@ -1,21 +1,49 @@
"""App Client API."""
from dbgpt.client.client import Client
from typing import List
from dbgpt.client.client import Client, ClientException
from dbgpt.client.schemas import AppModel
from dbgpt.serve.core import Result
async def get_app(client: Client, app_id: str):
async def get_app(client: Client, app_id: str) -> AppModel:
"""Get an app.
Args:
client (Client): The dbgpt client.
app_id (str): The app id.
Returns:
AppModel: The app model.
Raises:
ClientException: If the request failed.
"""
return await client.get("/apps/" + app_id)
try:
res = await client.get("/apps/" + app_id)
result: Result = res.json()
if result["success"]:
return AppModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to get app: {e}")
async def list_app(client: Client):
async def list_app(client: Client) -> List[AppModel]:
"""List apps.
Args:
client (Client): The dbgpt client.
Returns:
List[AppModel]: The list of app models.
Raises:
ClientException: If the request failed.
"""
return await client.get("/apps")
try:
res = await client.get("/apps")
result: Result = res.json()
if result["success"]:
return [AppModel(**app) for app in result["data"]["app_list"]]
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to list apps: {e}")

16
dbgpt/client/client.py
View File

@@ -24,17 +24,11 @@ class ClientException(Exception):
reason: Optional[str], the reason for the exception.
http_resp: Optional[httpx.Response], the HTTP response object.
"""
reason = json.loads(reason)
if http_resp:
self.status = http_resp.status_code
self.reason = http_resp.content
self.body = http_resp.content
self.headers = None
else:
self.status = status
self.reason = reason
self.body = None
self.headers = None
self.status = status
self.reason = reason
self.http_resp = http_resp
self.headers = http_resp.headers if http_resp else None
self.body = http_resp.text if http_resp else None
def __str__(self):
"""Return the error message."""

81
dbgpt/client/flow.py
View File

@@ -1,55 +1,114 @@
"""this module contains the flow client functions."""
from dbgpt.client.client import Client
from typing import List
from dbgpt.client.client import Client, ClientException
from dbgpt.core.awel.flow.flow_factory import FlowPanel
from dbgpt.serve.core import Result
async def create_flow(client: Client, flow: FlowPanel):
async def create_flow(client: Client, flow: FlowPanel) -> FlowPanel:
"""Create a new flow.
Args:
client (Client): The dbgpt client.
flow (FlowPanel): The flow panel.
"""
return await client.get("/awel/flows", flow.dict())
try:
res = await client.get("/awel/flows", flow.dict())
result: Result = res.json()
if result["success"]:
return FlowPanel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to create flow: {e}")
async def update_flow(client: Client, flow: FlowPanel):
async def update_flow(client: Client, flow: FlowPanel) -> FlowPanel:
"""Update a flow.
Args:
client (Client): The dbgpt client.
flow (FlowPanel): The flow panel.
Returns:
FlowPanel: The flow panel.
Raises:
ClientException: If the request failed.
"""
return await client.put("/awel/flows", flow.dict())
try:
res = await client.put("/awel/flows", flow.dict())
result: Result = res.json()
if result["success"]:
return FlowPanel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to update flow: {e}")
async def delete_flow(client: Client, flow_id: str):
async def delete_flow(client: Client, flow_id: str) -> FlowPanel:
"""
Delete a flow.
Args:
client (Client): The dbgpt client.
flow_id (str): The flow id.
Returns:
FlowPanel: The flow panel.
Raises:
ClientException: If the request failed.
"""
return await client.get("/awel/flows/" + flow_id)
try:
res = await client.delete("/awel/flows/" + flow_id)
result: Result = res.json()
if result["success"]:
return FlowPanel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to delete flow: {e}")
async def get_flow(client: Client, flow_id: str):
async def get_flow(client: Client, flow_id: str) -> FlowPanel:
"""
Get a flow.
Args:
client (Client): The dbgpt client.
flow_id (str): The flow id.
Returns:
FlowPanel: The flow panel.
Raises:
ClientException: If the request failed.
"""
return await client.get("/awel/flows/" + flow_id)
try:
res = await client.get("/awel/flows/" + flow_id)
result: Result = res.json()
if result["success"]:
return FlowPanel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to get flow: {e}")
async def list_flow(client: Client):
async def list_flow(client: Client) -> List[FlowPanel]:
"""
List flows.
Args:
client (Client): The dbgpt client.
Returns:
List[FlowPanel]: The list of flow panels.
Raises:
ClientException: If the request failed.
"""
return await client.get("/awel/flows")
try:
res = await client.get("/awel/flows")
result: Result = res.json()
if result["success"]:
return [FlowPanel(**flow) for flow in result["data"]["items"]]
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to list flows: {e}")

178
dbgpt/client/knowledge.py
View File

@@ -1,104 +1,220 @@
"""Knowledge API client."""
import json
from typing import List
from dbgpt.client.client import Client
from dbgpt.client.client import Client, ClientException
from dbgpt.client.schemas import DocumentModel, SpaceModel, SyncModel
from dbgpt.serve.core import Result
async def create_space(client: Client, app_model: SpaceModel):
async def create_space(client: Client, space_model: SpaceModel) -> SpaceModel:
"""Create a new space.
Args:
client (Client): The dbgpt client.
app_model (SpaceModel): The app model.
space_model (SpaceModel): The space model.
Returns:
SpaceModel: The space model.
Raises:
ClientException: If the request failed.
"""
return await client.post("/knowledge/spaces", app_model.dict())
try:
res = await client.post("/knowledge/spaces", space_model.dict())
result: Result = res.json()
if result["success"]:
return SpaceModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to create space: {e}")
async def update_space(client: Client, app_model: SpaceModel):
async def update_space(client: Client, space_model: SpaceModel) -> SpaceModel:
"""Update a document.
Args:
client (Client): The dbgpt client.
app_model (SpaceModel): The app model.
space_model (SpaceModel): The space model.
Returns:
SpaceModel: The space model.
Raises:
ClientException: If the request failed.
"""
return await client.put("/knowledge/spaces", app_model.dict())
try:
res = await client.put("/knowledge/spaces", space_model.dict())
result: Result = res.json()
if result["success"]:
return SpaceModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to update space: {e}")
async def delete_space(client: Client, space_id: str):
async def delete_space(client: Client, space_id: str) -> SpaceModel:
"""Delete a space.
Args:
client (Client): The dbgpt client.
app_id (str): The app id.
space_id (str): The space id.
Returns:
SpaceModel: The space model.
Raises:
ClientException: If the request failed.
"""
return await client.delete("/knowledge/spaces/" + space_id)
try:
res = await client.delete("/knowledge/spaces/" + space_id)
result: Result = res.json()
if result["success"]:
return SpaceModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to delete space: {e}")
async def get_space(client: Client, space_id: str):
async def get_space(client: Client, space_id: str) -> SpaceModel:
"""Get a document.
Args:
client (Client): The dbgpt client.
app_id (str): The app id.
space_id (str): The space id.
Returns:
SpaceModel: The space model.
Raises:
ClientException: If the request failed.
"""
return await client.get("/knowledge/spaces/" + space_id)
try:
res = await client.get("/knowledge/spaces/" + space_id)
result: Result = res.json()
if result["success"]:
return SpaceModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to get space: {e}")
async def list_space(client: Client):
"""List apps.
async def list_space(client: Client) -> List[SpaceModel]:
"""List spaces.
Args:
client (Client): The dbgpt client.
Returns:
List[SpaceModel]: The list of space models.
Raises:
ClientException: If the request failed.
"""
return await client.get("/knowledge/spaces")
try:
res = await client.get("/knowledge/spaces")
result: Result = res.json()
if result["success"]:
return [SpaceModel(**space) for space in result["data"]["items"]]
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to list spaces: {e}")
async def create_document(client: Client, doc_model: DocumentModel):
"""Create a new space.
async def create_document(client: Client, doc_model: DocumentModel) -> DocumentModel:
"""Create a new document.
Args:
client (Client): The dbgpt client.
doc_model (SpaceModel): The document model.
"""
return await client.post_param("/knowledge/documents", doc_model.dict())
try:
res = await client.post_param("/knowledge/documents", doc_model.dict())
result: Result = res.json()
if result["success"]:
return DocumentModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to create document: {e}")
async def delete_document(client: Client, document_id: str):
async def delete_document(client: Client, document_id: str) -> DocumentModel:
"""Delete a document.
Args:
client (Client): The dbgpt client.
app_id (str): The app id.
document_id (str): The document id.
Returns:
DocumentModel: The document model.
Raises:
ClientException: If the request failed.
"""
return await client.delete("/knowledge/documents/" + document_id)
try:
res = await client.delete("/knowledge/documents/" + document_id)
result: Result = res.json()
if result["success"]:
return DocumentModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to delete document: {e}")
async def get_document(client: Client, document_id: str):
async def get_document(client: Client, document_id: str) -> DocumentModel:
"""Get a document.
Args:
client (Client): The dbgpt client.
app_id (str): The app id.
document_id (str): The document id.
Returns:
DocumentModel: The document model.
Raises:
ClientException: If the request failed.
"""
return await client.get("/knowledge/documents/" + document_id)
try:
res = await client.get("/knowledge/documents/" + document_id)
result: Result = res.json()
if result["success"]:
return DocumentModel(**result["data"])
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to get document: {e}")
async def list_document(client: Client):
async def list_document(client: Client) -> List[DocumentModel]:
"""List documents.
Args:
client (Client): The dbgpt client.
"""
return await client.get("/knowledge/documents")
try:
res = await client.get("/knowledge/documents")
result: Result = res.json()
if result["success"]:
return [DocumentModel(**document) for document in result["data"]["items"]]
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to list documents: {e}")
async def sync_document(client: Client, sync_model: SyncModel):
async def sync_document(client: Client, sync_model: SyncModel) -> List:
"""Sync document.
Args:
client (Client): The dbgpt client.
sync_model (SyncModel): The sync model.
Returns:
List: The list of document ids.
Raises:
ClientException: If the request failed.
"""
return await client.post(
"/knowledge/documents/sync", [json.loads(sync_model.json())]
)
try:
res = await client.post(
"/knowledge/documents/sync", [json.loads(sync_model.json())]
)
result: Result = res.json()
if result["success"]:
return result["data"]
else:
raise ClientException(status=result["err_code"], reason=result)
except Exception as e:
raise ClientException(f"Failed to list documents: {e}")

10
dbgpt/client/schemas.py
View File

@@ -21,7 +21,7 @@ class ChatCompletionRequestBody(BaseModel):
messages: Union[str, List[str]] = Field(
..., description="User input messages", examples=["Hello", "How are you?"]
)
stream: bool = Field(default=False, description="Whether return stream")
stream: bool = Field(default=True, description="Whether return stream")
temperature: Optional[float] = Field(
default=None,
@@ -174,6 +174,10 @@ class AppModel(BaseModel):
class SpaceModel(BaseModel):
"""Space model."""
id: str = Field(
default=None,
description="space id",
)
name: str = Field(
default=None,
description="knowledge space name",
@@ -190,6 +194,10 @@ class SpaceModel(BaseModel):
default=None,
description="space owner",
)
context: Optional[str] = Field(
default=None,
description="space argument context",
)
class DocumentModel(BaseModel):

4
dbgpt/serve/flow/api/endpoints.py
View File

@@ -147,8 +147,8 @@ async def delete(uid: str, service: Service = Depends(get_service)) -> Result[No
Returns:
Result[None]: The response
"""
service.delete(uid)
return Result.succ(None)
inst = service.delete(uid)
return Result.succ(inst)
@router.get("/flows/{uid}")

12
dbgpt/serve/rag/api/schemas.py
View File

@@ -15,11 +15,17 @@ class SpaceServeRequest(BaseModel):
id: Optional[int] = Field(None, description="The space id")
name: str = Field(None, description="The space name")
"""vector_type: vector type"""
vector_type: str = Field(None, description="The vector type")
vector_type: str = Field("Chroma", description="The vector type")
"""desc: description"""
desc: str = Field(None, description="The description")
desc: Optional[str] = Field(None, description="The description")
"""owner: owner"""
owner: str = Field(None, description="The owner")
owner: Optional[str] = Field(None, description="The owner")
"""context: argument context"""
context: Optional[str] = Field(None, description="The context")
"""gmt_created: created time"""
gmt_created: Optional[str] = Field(None, description="The created time")
"""gmt_modified: modified time"""
gmt_modified: Optional[str] = Field(None, description="The modified time")
class DocumentServeRequest(BaseModel):

17
dbgpt/serve/rag/models/models.py
View File

@@ -38,7 +38,7 @@ class KnowledgeSpaceDao(BaseDao):
session.commit()
space_id = knowledge_space.id
session.close()
return space_id
return self.to_response(knowledge_space)
def get_knowledge_space(self, query: KnowledgeSpaceEntity):
"""Get knowledge space by query"""
@@ -81,11 +81,21 @@ class KnowledgeSpaceDao(BaseDao):
def update_knowledge_space(self, space: KnowledgeSpaceEntity):
"""Update knowledge space"""
session = self.get_raw_session()
session.merge(space)
request = SpaceServeRequest(id=space.id)
update_request = self.to_request(space)
query = self._create_query_object(session, request)
entry = query.first()
if entry is None:
raise Exception("Invalid request")
for key, value in update_request.dict().items(): # type: ignore
if value is not None:
setattr(entry, key, value)
session.merge(entry)
session.commit()
session.close()
return True
return self.to_response(space)
def delete_knowledge_space(self, space: KnowledgeSpaceEntity):
"""Delete knowledge space"""
@@ -127,6 +137,7 @@ class KnowledgeSpaceDao(BaseDao):
vector_type=entity.vector_type,
desc=entity.desc,
owner=entity.owner,
context=entity.context,
)
def to_response(self, entity: KnowledgeSpaceEntity) -> SpaceServeResponse:

4
dbgpt/serve/rag/service/service.py
View File

@@ -145,9 +145,7 @@ class Service(BaseService[KnowledgeSpaceEntity, SpaceServeRequest, SpaceServeRes
status_code=400,
detail=f"no space name named {request.name}",
)
space = spaces[0]
query_request = {"id": space.id}
update_obj = self._dao.update(query_request, update_request=request)
update_obj = self._dao.update_knowledge_space(self._dao.from_request(request))
return update_obj
async def create_document(