Refactor: Update llm type annotations to use BaseLanguageModel

Author: Eyobyb

src/sherpa_ai/actions/arxiv_search.py

@@ -1,5 +1,5 @@
-from typing import Any
-
+from typing import Any, Optional
+from langchain_core.language_models.base import BaseLanguageModel
 from sherpa_ai.actions.base import BaseRetrievalAction
 from sherpa_ai.tools import SearchArxivTool
 
@@ -46,7 +46,7 @@ class ArxivSearch(BaseRetrievalAction):
     """
     role_description: str
     task: str
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None
     description: str = SEARCH_SUMMARY_DESCRIPTION
     _search_tool: Any = None
 

src/sherpa_ai/actions/context_search.py

@@ -1,6 +1,7 @@
-from typing import Any
+from typing import Any, Optional
 
 from loguru import logger
+from langchain_core.language_models.base import BaseLanguageModel
 
 from sherpa_ai.actions.base import BaseRetrievalAction
 from sherpa_ai.connectors.vectorstores import get_vectordb
@@ -55,7 +56,7 @@ class ContextSearch(BaseRetrievalAction):
 
     role_description: str
     task: str
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None
     description: str = SEARCH_SUMMARY_DESCRIPTION
     _context: Any = None
 

src/sherpa_ai/actions/deliberation.py

@@ -1,5 +1,5 @@
-from typing import Any
-
+from typing import Any, Optional
+from langchain_core.language_models.base import BaseLanguageModel
 from sherpa_ai.actions.base import BaseAction
 
 
@@ -50,7 +50,7 @@ class Deliberation(BaseAction):
     """
     # TODO: Make a version of Deliberation action that considers the context
     role_description: str
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None
     description: str = DELIBERATION_DESCRIPTION
 
     # Override the name and args from BaseAction

src/sherpa_ai/actions/google_search.py

@@ -1,7 +1,7 @@
-from typing import Any
+from typing import Any, Optional
 
 from loguru import logger
-
+from langchain_core.language_models.base import BaseLanguageModel
 from sherpa_ai.actions.base import BaseRetrievalAction
 from sherpa_ai.config.task_config import AgentConfig
 from sherpa_ai.tools import SearchTool
@@ -65,7 +65,7 @@ class GoogleSearch(BaseRetrievalAction):
     """
     role_description: str
     task: str
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None
     description: str = SEARCH_SUMMARY_DESCRIPTION
     config: AgentConfig = AgentConfig()
     _search_tool: Any = None

src/sherpa_ai/actions/planning.py

@@ -1,7 +1,7 @@
 from typing import Any, Optional
 
 from loguru import logger
-
+from langchain_core.language_models.base import BaseLanguageModel
 from sherpa_ai.actions.base import BaseAction
 
 
@@ -216,7 +216,7 @@ class TaskPlanning(BaseAction):
         Task: Summarize the findings about quantum computing
     """
 
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None
     num_steps: int = 5
     prompt: str = PLANNING_PROMPT
     revision_prompt: str = REVISION_PROMPT

src/sherpa_ai/actions/synthesize.py

@@ -1,6 +1,6 @@
-from typing import Any
+from typing import Any, Optional
 
-from langchain_core.language_models import BaseLanguageModel 
+from langchain_core.language_models.base import BaseLanguageModel 
 from loguru import logger 
 
 from sherpa_ai.actions.base import BaseAction
@@ -42,7 +42,7 @@ class SynthesizeOutput(BaseAction):
     """
 
     role_description: str
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None
     description: str = None
     add_citation: bool = False
 

src/sherpa_ai/actions/utils/refinement.py

@@ -3,12 +3,13 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 import numpy as np
 from nltk import tokenize
 from numpy.typing import ArrayLike
 from pydantic import BaseModel
+from langchain_core.language_models.base import BaseLanguageModel
 
 
 SEARCH_SUMMARY_DESCRIPTION = """Question：{question} 
@@ -65,7 +66,7 @@ class RefinementByQuery(BaseRefinement):
         >>> refinement = RefinementByQuery(llm=my_llm)
         >>> results = refinement.refinement(documents, query)
     """
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None 
     description: str = SEARCH_SUMMARY_DESCRIPTION
     k: int = 3
 
@@ -108,7 +109,7 @@ class RefinementBySentence(BaseRefinement):
         >>> refinement = RefinementBySentence(llm=my_llm)
         >>> results = refinement.refinement(documents, query)
     """
-    llm: Any = None  # The BaseLanguageModel from LangChain is not compatible with Pydantic 2 yet
+    llm: Optional[BaseLanguageModel] = None 
     description: str = SEARCH_SUMMARY_DESCRIPTION_SENT
 
     def refinement(self, documents: list[str], query: str) -> list[str]:

src/sherpa_ai/agents/base.py

@@ -8,6 +8,7 @@
 
 from loguru import logger
 from pydantic import BaseModel, ConfigDict
+from langchain_core.language_models.base import BaseLanguageModel
 
 from sherpa_ai.actions.base import BaseAction
 from sherpa_ai.actions.exceptions import (
@@ -73,7 +74,7 @@ class BaseAgent(ABC, BaseModel):
     validations: List[BaseOutputProcessor] = []
     feedback_agent_name: str = "critic"
     global_regen_max: int = 12
-    llm: Any = None
+    llm: Optional[BaseLanguageModel] = None
     prompt_template: PromptTemplate = None
 
     # Checks whether the execution of the agent should be stopped

src/sherpa_ai/policies/react_policy.py

@@ -12,6 +12,7 @@
 
 from loguru import logger
 from pydantic import BaseModel, ConfigDict
+from langchain_core.language_models.base import BaseLanguageModel
 
 from sherpa_ai.policies.base import BasePolicy, PolicyOutput
 from sherpa_ai.policies.exceptions import SherpaPolicyException
@@ -55,8 +56,7 @@ class ReactPolicy(BasePolicy):
 
     role_description: str
     output_instruction: str
-    # Cannot use langchain's BaseLanguageModel due to they are using Pydantic v1
-    llm: Any = None
+    llm: Optional[BaseLanguageModel] = None
     prompt_template: PromptTemplate = PromptTemplate("./sherpa_ai/prompts/prompts.json")
     response_format: dict = {
         "command": {

src/sherpa_ai/policies/react_sm_policy.py

@@ -13,6 +13,7 @@
 
 from loguru import logger
 from pydantic import BaseModel, ConfigDict
+from langchain_core.language_models.base import BaseLanguageModel
 
 from sherpa_ai.actions.base import BaseAction
 from sherpa_ai.policies.base import BasePolicy, PolicyOutput
@@ -55,8 +56,7 @@ class ReactStateMachinePolicy(BasePolicy):
     )
     role_description: str
     output_instruction: str
-    # Cannot use langchain's BaseLanguageModel due to they are using Pydantic v1
-    llm: Any = None
+    llm: Optional[BaseLanguageModel] = None
     prompt_template: PromptTemplate = PromptTemplate("./sherpa_ai/prompts/prompts.json")
 
     response_format: dict = {