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

dbgpt/serve/rag/api/schemas.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 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
 
 from ..config import SERVE_APP_NAME_HUMP
@@ -29,71 +29,102 @@ class SpaceServeRequest(BaseModel):
 
 
 class DocumentServeRequest(BaseModel):
-    id: int = Field(None, description="The doc id")
-    doc_name: str = Field(None, description="doc name")
+    id: Optional[int] = Field(None, description="The doc id")
+    doc_name: Optional[str] = Field(None, description="doc name")
     """doc_type: document type"""
-    doc_type: str = Field(None, description="The doc type")
+    doc_type: Optional[str] = Field(None, description="The doc type")
     """content: description"""
-    content: str = Field(None, description="content")
+    content: Optional[str] = Field(None, description="content")
     """doc file"""
     doc_file: UploadFile = File(...)
     """doc_source: doc source"""
-    doc_source: str = None
+    doc_source: Optional[str] = Field(None, description="doc source")
     """doc_source: doc source"""
-    space_id: str = None
+    space_id: Optional[str] = Field(None, description="space id")
 
 
 class DocumentServeResponse(BaseModel):
-    id: int = Field(None, description="The doc id")
-    doc_name: str = Field(None, description="doc type")
+    id: Optional[int] = Field(None, description="The doc id")
+    doc_name: Optional[str] = Field(None, description="doc type")
     """vector_type: vector type"""
-    doc_type: str = Field(None, description="The doc content")
+    doc_type: Optional[str] = Field(None, description="The doc content")
     """desc: description"""
-    content: str = Field(None, description="content")
+    content: Optional[str] = Field(None, description="content")
     """vector ids"""
-    vector_ids: str = Field(None, description="vector ids")
+    vector_ids: Optional[str] = Field(None, description="vector ids")
     """doc_source: doc source"""
-    doc_source: str = None
+    doc_source: Optional[str] = Field(None, description="doc source")
     """doc_source: doc source"""
-    space: str = None
+    space: Optional[str] = Field(None, description="space name")
 
 
 class KnowledgeSyncRequest(BaseModel):
     """Sync request"""
 
     """doc_ids: doc ids"""
-    doc_id: int = Field(None, description="The doc id")
+    doc_id: Optional[int] = Field(None, description="The doc id")
 
     """space id"""
-    space_id: str = Field(None, description="space id")
+    space_id: Optional[str] = Field(None, description="space id")
 
     """model_name: model name"""
     model_name: Optional[str] = Field(None, description="model name")
 
     """chunk_parameters: chunk parameters 
     """
-    chunk_parameters: ChunkParameters = Field(None, description="chunk parameters")
+    chunk_parameters: Optional[ChunkParameters] = Field(
+        None, description="chunk parameters"
+    )
 
 
 class SpaceServeResponse(BaseModel):
     """Flow response model"""
 
+    model_config = ConfigDict(title=f"ServeResponse for {SERVE_APP_NAME_HUMP}")
+
     """name: knowledge space name"""
 
     """vector_type: vector type"""
-    id: int = Field(None, description="The space id")
-    name: str = Field(None, description="The space name")
+    id: Optional[int] = Field(None, description="The space id")
+    name: Optional[str] = Field(None, description="The space name")
     """vector_type: vector type"""
-    vector_type: str = Field(None, description="The vector type")
+    vector_type: Optional[str] = Field(None, description="The vector type")
     """desc: description"""
-    desc: str = Field(None, description="The description")
+    desc: Optional[str] = Field(None, description="The description")
     """context: argument context"""
-    context: str = Field(None, description="The context")
+    context: Optional[str] = Field(None, description="The context")
     """owner: owner"""
-    owner: str = Field(None, description="The owner")
+    owner: Optional[str] = Field(None, description="The owner")
     """sys code"""
-    sys_code: str = Field(None, description="The sys code")
+    sys_code: Optional[str] = Field(None, description="The sys code")
 
     # TODO define your own fields here
-    class Config:
-        title = f"ServerResponse for {SERVE_APP_NAME_HUMP}"
+
+
+class DocumentChunkVO(BaseModel):
+    id: int = Field(..., description="document chunk id")
+    document_id: int = Field(..., description="document id")
+    doc_name: str = Field(..., description="document name")
+    doc_type: str = Field(..., description="document type")
+    content: str = Field(..., description="document content")
+    meta_info: str = Field(..., description="document meta info")
+    gmt_created: str = Field(..., description="document create time")
+    gmt_modified: str = Field(..., description="document modify time")
+
+
+class DocumentVO(BaseModel):
+    """Document Entity."""
+
+    id: int = Field(..., description="document id")
+    doc_name: str = Field(..., description="document name")
+    doc_type: str = Field(..., description="document type")
+    space: str = Field(..., description="document space name")
+    chunk_size: int = Field(..., description="document chunk size")
+    status: str = Field(..., description="document status")
+    last_sync: str = Field(..., description="document last sync time")
+    content: str = Field(..., description="document content")
+    result: Optional[str] = Field(None, description="document result")
+    vector_ids: Optional[str] = Field(None, description="document vector ids")
+    summary: Optional[str] = Field(None, description="document summary")
+    gmt_created: str = Field(..., description="document create time")
+    gmt_modified: str = Field(..., description="document modify time")

dbgpt/serve/rag/models/models.py

@@ -3,6 +3,7 @@ from typing import Any, Dict, List, Union
 
 from sqlalchemy import Column, DateTime, Integer, String, Text
 
+from dbgpt._private.pydantic import model_to_dict
 from dbgpt.serve.rag.api.schemas import SpaceServeRequest, SpaceServeResponse
 from dbgpt.storage.metadata import BaseDao, Model
 
@@ -89,7 +90,7 @@ class KnowledgeSpaceDao(BaseDao):
         entry = query.first()
         if entry is None:
             raise Exception("Invalid request")
-        for key, value in update_request.dict().items():  # type: ignore
+        for key, value in model_to_dict(update_request).items():  # type: ignore
             if value is not None:
                 setattr(entry, key, value)
         session.merge(entry)
@@ -117,7 +118,9 @@ class KnowledgeSpaceDao(BaseDao):
             T: The entity
         """
         request_dict = (
-            request.dict() if isinstance(request, SpaceServeRequest) else request
+            model_to_dict(request)
+            if isinstance(request, SpaceServeRequest)
+            else request
         )
         entity = KnowledgeSpaceEntity(**request_dict)
         return entity

dbgpt/serve/rag/service/service.py

@@ -40,6 +40,7 @@ from dbgpt.util.tracer import root_tracer, trace
 from ..api.schemas import (
     DocumentServeRequest,
     DocumentServeResponse,
+    DocumentVO,
     KnowledgeSyncRequest,
     SpaceServeRequest,
     SpaceServeResponse,
@@ -419,7 +420,7 @@ class Service(BaseService[KnowledgeSpaceEntity, SpaceServeRequest, SpaceServeRes
     def _sync_knowledge_document(
         self,
         space_id,
-        doc: KnowledgeDocumentEntity,
+        doc_vo: DocumentVO,
         chunk_parameters: ChunkParameters,
     ) -> List[Chunk]:
         """sync knowledge document chunk into vector store"""
@@ -431,6 +432,8 @@ class Service(BaseService[KnowledgeSpaceEntity, SpaceServeRequest, SpaceServeRes
         )
         from dbgpt.storage.vector_store.base import VectorStoreConfig
 
+        doc = KnowledgeDocumentEntity.from_document_vo(doc_vo)
+
         space = self.get({"id": space_id})
         config = VectorStoreConfig(
             name=space.name,