Merge pull request #413 from Aggregate-Intellect/fix/pydantic-warning

20001LastOrder · web-flow · commit d4ea413c9178 · 2024-09-13T20:37:01.000-04:00
upgrade pydantic from v1 to v2
diff --git a/src/poetry.lock b/src/poetry.lock
diff --git a/src/sherpa_ai/actions/arxiv_search.py b/src/sherpa_ai/actions/arxiv_search.py
@@ -19,9 +19,9 @@
 class ArxivSearch(BaseRetrievalAction):
     role_description: str
     task: str
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     description: str = SEARCH_SUMMARY_DESCRIPTION
-    _search_tool: Any
+    _search_tool: Any = None
 
     # Override the name and args from BaseAction
     name: str = "ArxivSearch"
diff --git a/src/sherpa_ai/actions/belief_actions.py b/src/sherpa_ai/actions/belief_actions.py
@@ -6,11 +6,11 @@
 
 from sherpa_ai.actions.base import BaseAction
 from sherpa_ai.memory.belief import Belief
+from pydantic import ConfigDict
 
 
 class UpdateBelief(BaseAction):
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     name: str = "update_belief"
     args: dict = {
@@ -39,8 +39,7 @@ def execute(self, key: str, value: str) -> str:
 
 
 class RetrieveBelief(BaseAction):
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     name: str = "retrieve_belief"
     args: dict = {
diff --git a/src/sherpa_ai/actions/context_search.py b/src/sherpa_ai/actions/context_search.py
@@ -22,9 +22,9 @@
 class ContextSearch(BaseRetrievalAction):
     role_description: str
     task: str
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     description: str = SEARCH_SUMMARY_DESCRIPTION
-    _context: Any
+    _context: Any = None
 
     # Override the name and args from BaseAction
     name: str = "Context Search"
diff --git a/src/sherpa_ai/actions/deliberation.py b/src/sherpa_ai/actions/deliberation.py
@@ -20,7 +20,7 @@
 class Deliberation(BaseAction):
     # TODO: Make a version of Deliberation action that considers the context
     role_description: str
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     description: str = DELIBERATION_DESCRIPTION
 
     # Override the name and args from BaseAction
diff --git a/src/sherpa_ai/actions/google_search.py b/src/sherpa_ai/actions/google_search.py
@@ -34,10 +34,10 @@
 class GoogleSearch(BaseRetrievalAction):
     role_description: str
     task: str
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     description: str = SEARCH_SUMMARY_DESCRIPTION
     config: AgentConfig = AgentConfig()
-    _search_tool: Any
+    _search_tool: Any = None
 
     # Override the name and args from BaseAction
     name: str = "Google Search"
diff --git a/src/sherpa_ai/actions/planning.py b/src/sherpa_ai/actions/planning.py
@@ -107,7 +107,7 @@ def from_dict(cls, data):
 
 
 class TaskPlanning(BaseAction):
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     num_steps: int = 5
     prompt: str = PLANNING_PROMPT
     revision_prompt: str = REVISION_PROMPT
diff --git a/src/sherpa_ai/actions/synthesize.py b/src/sherpa_ai/actions/synthesize.py
@@ -33,7 +33,7 @@
 
 class SynthesizeOutput(BaseAction):
     role_description: str
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     description: str = SYNTHESIZE_DESCRIPTION
     add_citation: bool = False
 
diff --git a/src/sherpa_ai/actions/utils/refinement.py b/src/sherpa_ai/actions/utils/refinement.py
@@ -35,7 +35,7 @@ def refinement(self, documents: list[str], **kwargs) -> str:
 
 
 class RefinementByQuery(BaseRefinement):
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     description: str = SEARCH_SUMMARY_DESCRIPTION
     k: int = 3
 
@@ -51,7 +51,7 @@ def refinement(self, documents: list[str], query: str) -> list[str]:
 
 
 class RefinementBySentence(BaseRefinement):
-    llm: Any  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
     description: str = SEARCH_SUMMARY_DESCRIPTION_SENT
 
     def refinement(self, documents: list[str], query: str) -> list[str]:
diff --git a/src/sherpa_ai/actions/utils/reranking.py b/src/sherpa_ai/actions/utils/reranking.py
@@ -21,7 +21,7 @@ def rerank(self, documents: list[str], **kwargs) -> str:
 
 
 class RerankingByQuery(BaseReranking):
-    embeddings: Any  # takes an Embedding Object from LangChain, use Any since it is not compatible with Pydantic 2 yet
+    embeddings: Any = None  # takes an Embedding Object from LangChain, use Any since it is not compatible with Pydantic 2 yet
     distance_metric: Callable[[ArrayLike, ArrayLike], float] = cosine_similarity
 
     def rerank(self, documents: list[str], query: str) -> str:
diff --git a/src/sherpa_ai/config/task_config.py b/src/sherpa_ai/config/task_config.py
@@ -4,7 +4,7 @@
 from typing import List, Optional, Tuple
 from urllib.parse import urlparse
 
-from pydantic import BaseModel, computed_field, validator
+from pydantic import field_validator, BaseModel, computed_field
 
 
 class AgentConfig(BaseModel):
@@ -13,7 +13,8 @@ class AgentConfig(BaseModel):
     do_reflect: bool = False
     use_task_agent: bool = False
 
-    @validator("gsite", pre=True)
+    @field_validator("gsite", mode="before")
+    @classmethod
     def parse_gsite(cls, value: Optional[str]) -> list[str]:
         if value is None:
             return []
diff --git a/src/sherpa_ai/policies/agent_feedback_policy.py b/src/sherpa_ai/policies/agent_feedback_policy.py
@@ -4,6 +4,7 @@
 from sherpa_ai.events import Event, EventType
 from sherpa_ai.memory import Belief
 from sherpa_ai.policies import ReactPolicy
+from pydantic import ConfigDict
 
 AGENT_FEEDBACK_DESCRIPTION = """You are an intelligent assistant helping the user to complete their task. You have the following task to complete:
 {task}
@@ -25,9 +26,7 @@ class AgentFeedbackPolicy(ReactPolicy):
     """
     Select the best next action based on the feedback from an agent
     """
-
-    class Config:
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     # Prompt used to generate question to the user
     agent_feedback_description: str = AGENT_FEEDBACK_DESCRIPTION
diff --git a/src/sherpa_ai/policies/base.py b/src/sherpa_ai/policies/base.py
@@ -21,7 +21,7 @@ class PolicyOutput(BaseModel):
         args (dict): The arguments for the selected action
     """
 
-    action: Any  # TODO: Currently, we cannot use BaseAction since it does not inherit from Base class.
+    action: Any = None  # TODO: Currently, we cannot use BaseAction since it does not inherit from Base class.
     args: dict
 
 
diff --git a/src/sherpa_ai/policies/react_policy.py b/src/sherpa_ai/policies/react_policy.py
@@ -52,7 +52,7 @@ class ReactPolicy(BasePolicy):
 
     role_description: str
     output_instruction: str
-    llm: Any  # Cannot use langchain's BaseLanguageModel due to they are using Pydantic v1
+    llm: Any = None  # Cannot use langchain's BaseLanguageModel due to they are using Pydantic v1
     description: str = SELECTION_DESCRIPTION
     response_format: dict = {
         "command": {
