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

dbgpt/agent/core/agent_manage.py

@@ -112,7 +112,7 @@ class AgentManager(BaseComponent):
         """Return the description of an agent by name."""
         return self._agents[name][1].desc
 
-    def all_agents(self):
+    def all_agents(self) -> Dict[str, str]:
         """Return a dictionary of all registered agents and their descriptions."""
         result = {}
         for name, value in self._agents.items():

dbgpt/agent/core/base_agent.py

@@ -5,7 +5,7 @@ import json
 import logging
 from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
 
-from dbgpt._private.pydantic import Field
+from dbgpt._private.pydantic import ConfigDict, Field
 from dbgpt.core import LLMClient, ModelMessageRoleType
 from dbgpt.util.error_types import LLMChatError
 from dbgpt.util.tracer import SpanType, root_tracer
@@ -27,6 +27,8 @@ logger = logging.getLogger(__name__)
 class ConversableAgent(Role, Agent):
     """ConversableAgent is an agent that can communicate with other agents."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     agent_context: Optional[AgentContext] = Field(None, description="Agent context")
     actions: List[Action] = Field(default_factory=list)
     resources: List[AgentResource] = Field(default_factory=list)
@@ -38,11 +40,6 @@ class ConversableAgent(Role, Agent):
     llm_client: Optional[AIWrapper] = None
     oai_system_message: List[Dict] = Field(default_factory=list)
 
-    class Config:
-        """Pydantic configuration."""
-
-        arbitrary_types_allowed = True
-
     def __init__(self, **kwargs):
         """Create a new agent."""
         Role.__init__(self, **kwargs)
@@ -377,8 +374,10 @@ class ConversableAgent(Role, Agent):
                         **act_extent_param,
                     )
                     if act_out:
-                        reply_message.action_report = act_out.dict()
-                    span.metadata["action_report"] = act_out.dict() if act_out else None
+                        reply_message.action_report = act_out.to_dict()
+                    span.metadata["action_report"] = (
+                        act_out.to_dict() if act_out else None
+                    )
 
                 with root_tracer.start_span(
                     "agent.generate_reply.verify",
@@ -496,7 +495,7 @@ class ConversableAgent(Role, Agent):
                     "recipient": self.get_name(),
                     "reviewer": reviewer.get_name() if reviewer else None,
                     "need_resource": need_resource.to_dict() if need_resource else None,
-                    "rely_action_out": last_out.dict() if last_out else None,
+                    "rely_action_out": last_out.to_dict() if last_out else None,
                     "conv_uid": self.not_null_agent_context.conv_id,
                     "action_index": i,
                     "total_action": len(self.actions),
@@ -508,7 +507,7 @@ class ConversableAgent(Role, Agent):
                     rely_action_out=last_out,
                     **kwargs,
                 )
-                span.metadata["action_out"] = last_out.dict() if last_out else None
+                span.metadata["action_out"] = last_out.to_dict() if last_out else None
         return last_out
 
     async def correctness_check(

dbgpt/agent/core/base_team.py

@@ -3,7 +3,7 @@
 import logging
 from typing import Dict, List, Optional, Tuple, Union
 
-from dbgpt._private.pydantic import BaseModel, Field
+from dbgpt._private.pydantic import BaseModel, ConfigDict, Field
 
 from ..actions.action import ActionOutput
 from .agent import Agent, AgentMessage
@@ -68,16 +68,13 @@ def _content_str(content: Union[str, List, None]) -> str:
 class Team(BaseModel):
     """Team class for managing a group of agents in a team chat."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     agents: List[Agent] = Field(default_factory=list)
     messages: List[Dict] = Field(default_factory=list)
     max_round: int = 100
     is_team: bool = True
 
-    class Config:
-        """Pydantic model configuration."""
-
-        arbitrary_types_allowed = True
-
     def __init__(self, **kwargs):
         """Create a new Team instance."""
         super().__init__(**kwargs)
@@ -122,6 +119,8 @@ class Team(BaseModel):
 class ManagerAgent(ConversableAgent, Team):
     """Manager Agent class."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     profile: str = "TeamManager"
     goal: str = "manage all hired intelligent agents to complete mission objectives"
     constraints: List[str] = []
@@ -132,11 +131,6 @@ class ManagerAgent(ConversableAgent, Team):
     # of the agent has already been retried.
     max_retry_count: int = 1
 
-    class Config:
-        """Pydantic model configuration."""
-
-        arbitrary_types_allowed = True
-
     def __init__(self, **kwargs):
         """Create a new ManagerAgent instance."""
         ConversableAgent.__init__(self, **kwargs)

dbgpt/agent/core/llm/llm.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from enum import Enum
 from typing import Any, Dict, List, Optional, Type
 
-from dbgpt._private.pydantic import BaseModel, Field
+from dbgpt._private.pydantic import BaseModel, ConfigDict, Field
 from dbgpt.core import LLMClient, ModelMetadata, ModelRequest
 
 logger = logging.getLogger(__name__)
@@ -106,11 +106,8 @@ def register_llm_strategy(
 class LLMConfig(BaseModel):
     """LLM configuration."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     llm_client: Optional[LLMClient] = Field(default_factory=LLMClient)
     llm_strategy: LLMStrategyType = Field(default=LLMStrategyType.Default)
     strategy_context: Optional[Any] = None
-
-    class Config:
-        """Pydantic model config."""
-
-        arbitrary_types_allowed = True

dbgpt/agent/core/role.py

@@ -2,33 +2,30 @@
 from abc import ABC
 from typing import List, Optional
 
-from dbgpt._private.pydantic import BaseModel
+from dbgpt._private.pydantic import BaseModel, ConfigDict, Field
 
 
 class Role(ABC, BaseModel):
     """Role class for role-based conversation."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     profile: str = ""
     name: str = ""
-    resource_introduction = ""
+    resource_introduction: str = ""
     goal: str = ""
 
     expand_prompt: str = ""
 
-    fixed_subgoal: Optional[str] = None
+    fixed_subgoal: Optional[str] = Field(None, description="Fixed subgoal")
 
-    constraints: List[str] = []
+    constraints: List[str] = Field(default_factory=list, description="Constraints")
     examples: str = ""
     desc: str = ""
     language: str = "en"
     is_human: bool = False
     is_team: bool = False
 
-    class Config:
-        """Pydantic config."""
-
-        arbitrary_types_allowed = True
-
     def prompt_template(
         self,
         specified_prompt: Optional[str] = None,

dbgpt/agent/core/user_proxy_agent.py

@@ -8,7 +8,7 @@ class UserProxyAgent(ConversableAgent):
     That can execute code and provide feedback to the other agents.
     """
 
-    name = "User"
+    name: str = "User"
     profile: str = "Human"
 
     desc: str = (
@@ -16,4 +16,4 @@ class UserProxyAgent(ConversableAgent):
         "Plan execution needs to be approved by this admin."
     )
 
-    is_human = True
+    is_human: bool = True