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

dbgpt/client/client.py

@@ -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:

dbgpt/client/datasource.py

@@ -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"])

dbgpt/client/flow.py

@@ -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"])

dbgpt/client/knowledge.py

@@ -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"]:

dbgpt/client/schema.py

@@ -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")