feat(core): Upgrade pydantic to 2.x (#1428)

This commit is contained in:
Fangyin Cheng
2024-04-20 09:41:16 +08:00
committed by GitHub
parent baa1e3f9f6
commit 57be1ece18
103 changed files with 1146 additions and 534 deletions

View File

@@ -7,6 +7,7 @@ from urllib.parse import urlparse
import httpx
from dbgpt._private.pydantic import model_to_dict
from dbgpt.core.schema.api import ChatCompletionResponse, ChatCompletionStreamResponse
from .schema import ChatCompletionRequestBody
@@ -167,7 +168,7 @@ class Client:
enable_vis=enable_vis,
)
response = await self._http_client.post(
self._api_url + "/chat/completions", json=request.dict()
self._api_url + "/chat/completions", json=model_to_dict(request)
)
if response.status_code == 200:
json_data = json.loads(response.text)
@@ -242,7 +243,7 @@ class Client:
incremental=incremental,
enable_vis=enable_vis,
)
async for chat_completion_response in self._chat_stream(request.dict()):
async for chat_completion_response in self._chat_stream(model_to_dict(request)):
yield chat_completion_response
async def _chat_stream(
@@ -262,6 +263,7 @@ class Client:
headers={},
) as response:
if response.status_code == 200:
sse_data = ""
async for line in response.aiter_lines():
try:
if line.strip() == "data: [DONE]":
@@ -277,7 +279,9 @@ class Client:
)
yield chat_completion_response
except Exception as e:
raise e
raise Exception(
f"Failed to parse SSE data: {e}, sse_data: {sse_data}"
)
else:
try:

View File

@@ -1,6 +1,7 @@
"""this module contains the datasource client functions."""
from typing import List
from dbgpt._private.pydantic import model_to_dict
from dbgpt.core.schema.api import Result
from .client import Client, ClientException
@@ -17,7 +18,7 @@ async def create_datasource(
datasource (DatasourceModel): The datasource model.
"""
try:
res = await client.get("/datasources", datasource.dict())
res = await client.get("/datasources", model_to_dict(datasource))
result: Result = res.json()
if result["success"]:
return DatasourceModel(**result["data"])
@@ -41,7 +42,7 @@ async def update_datasource(
ClientException: If the request failed.
"""
try:
res = await client.put("/datasources", datasource.dict())
res = await client.put("/datasources", model_to_dict(datasource))
result: Result = res.json()
if result["success"]:
return DatasourceModel(**result["data"])

View File

@@ -15,7 +15,7 @@ async def create_flow(client: Client, flow: FlowPanel) -> FlowPanel:
flow (FlowPanel): The flow panel.
"""
try:
res = await client.get("/awel/flows", flow.dict())
res = await client.get("/awel/flows", flow.to_dict())
result: Result = res.json()
if result["success"]:
return FlowPanel(**result["data"])
@@ -37,7 +37,7 @@ async def update_flow(client: Client, flow: FlowPanel) -> FlowPanel:
ClientException: If the request failed.
"""
try:
res = await client.put("/awel/flows", flow.dict())
res = await client.put("/awel/flows", flow.to_dict())
result: Result = res.json()
if result["success"]:
return FlowPanel(**result["data"])

View File

@@ -2,6 +2,7 @@
import json
from typing import List
from dbgpt._private.pydantic import model_to_dict, model_to_json
from dbgpt.core.schema.api import Result
from .client import Client, ClientException
@@ -20,7 +21,7 @@ async def create_space(client: Client, space_model: SpaceModel) -> SpaceModel:
ClientException: If the request failed.
"""
try:
res = await client.post("/knowledge/spaces", space_model.dict())
res = await client.post("/knowledge/spaces", model_to_dict(space_model))
result: Result = res.json()
if result["success"]:
return SpaceModel(**result["data"])
@@ -42,7 +43,7 @@ async def update_space(client: Client, space_model: SpaceModel) -> SpaceModel:
ClientException: If the request failed.
"""
try:
res = await client.put("/knowledge/spaces", space_model.dict())
res = await client.put("/knowledge/spaces", model_to_dict(space_model))
result: Result = res.json()
if result["success"]:
return SpaceModel(**result["data"])
@@ -126,7 +127,7 @@ async def create_document(client: Client, doc_model: DocumentModel) -> DocumentM
"""
try:
res = await client.post_param("/knowledge/documents", doc_model.dict())
res = await client.post_param("/knowledge/documents", model_to_dict(doc_model))
result: Result = res.json()
if result["success"]:
return DocumentModel(**result["data"])
@@ -210,7 +211,7 @@ async def sync_document(client: Client, sync_model: SyncModel) -> List:
"""
try:
res = await client.post(
"/knowledge/documents/sync", [json.loads(sync_model.json())]
"/knowledge/documents/sync", [json.loads(model_to_json(sync_model))]
)
result: Result = res.json()
if result["success"]:

View File

@@ -6,7 +6,7 @@ from typing import Any, Dict, List, Optional, Union
from fastapi import File, UploadFile
from dbgpt._private.pydantic import BaseModel, Field
from dbgpt._private.pydantic import BaseModel, ConfigDict, Field
from dbgpt.rag.chunk_manager import ChunkParameters
@@ -60,7 +60,7 @@ class ChatCompletionRequestBody(BaseModel):
"or in full each time. "
"If this parameter is not provided, the default is full return.",
)
enable_vis: str = Field(
enable_vis: bool = Field(
default=True, description="response content whether to output vis label"
)
@@ -267,6 +267,8 @@ class DocumentModel(BaseModel):
class SyncModel(BaseModel):
"""Sync model."""
model_config = ConfigDict(protected_namespaces=())
"""doc_id: doc id"""
doc_id: str = Field(None, description="The doc id")